The Tuned Lens replaces the logit lens with a trained per-layer affine correction

Belrose et al. replace the logit lens's direct reuse of the final unembedding on intermediate states with a trained per-layer affine translator, distilled via KL divergence to the model's own final output [belrose-etal-2023-tuned-lens] This removes two documented logit-lens failure modes: a systematic 4-5 bit-per-layer vocabulary bias and unreliability from hidden-state covariance drifting across layers [belrose-etal-2023-tuned-lens] A causal-basis extraction finds the lens's most-sensitive directions strongly match the model's real causal directions (Spearman rho=0.89), with none influential to the lens but not the model [belrose-etal-2023-tuned-lens] Tested across GPT-2, GPT-Neo, GPT-NeoX-20B, OPT, BLOOM, and the full Pythia suite, it gives lower perplexity than the logit lens, transfers to a fine-tune (Vicuna-13B from LLaMA-13B), and detects prompt injection at near-perfect AUROC [belrose-etal-2023-tuned-lens]