Logit lens

Overview

A transformer's hidden states can be read as successive refinements of a guess for the next token 2303.08112 §1 2303.08112 §2. The problem is that the last layer's decoder does not transfer cleanly to earlier layers: the intermediate states can sit in a different basis from the final output, so a direct read-out can be biased 2303.08112 §3. Every variant below decodes a hidden state at layer l into a distribution over the vocabulary and compares it with what the model finally predicts. The shared invariant is that the base model is frozen; only the read-out changes.

Variants

Logit lens (zero residual)

Apply the model's final layer norm and unembedding matrix directly to the hidden state. This treats the remaining residual updates from layer l to the output as zero, which needs no training 2303.08112 §2 2303.08112 eq. 3. The paper reports that for BLOOM and OPT 125M the top probability often falls on the input token rather than on a plausible continuation, in more than half of the layers 2303.08112 §2 2303.08112 app. B.1. For GPT-Neo-2.7B the bias against the final output is about 4 to 5 bits for most layers 2303.08112 §2 2303.08112 fig. 3. Card: Eliciting Latent Predictions from Transformers with the Tuned Lens.

Logit lens with the last update kept

A variant keeps the final layer's residual update instead of dropping it 2303.08112 eq. 4 2303.08112 §2. The paper reports that it works better for some models and not for others 2303.08112 fig. 1 2303.08112 §2. No model-independent rule is given.

Tuned lens (per-layer affine translator)

Train one affine map A_l h + b_l for each layer, then apply the model's own final norm and unembedding 2303.08112 §3 2303.08112 eq. 8. The map is d x d plus a bias of size d, so it is far smaller than a new vocabulary-sized read-out 2303.08112 §3. Each map is trained by distillation: the loss is the KL divergence from the final-layer distribution to the lens output, averaged over positions 2303.08112 §3 2303.08112 eq. 9. The final-layer distribution is a soft target, so the map is not trained on ground-truth labels and is not meant to learn facts beyond what the model already encodes 2303.08112 §3. Reported results: lower perplexity than the logit lens across the models tested, with or without the final layer in the probe 2303.08112 §3 2303.08112 fig. 4 2303.08112 app. A. The paper also reports that training on one model can transfer to its fine-tune with only a small loss in fit: at worst a 0.3 bits-per-byte increase in KL on RedPajama, and no significant difference on the helpful-harmless set 2303.08112 §3 2303.08112 app. A.

Applications built on the lens

These are uses reported in the same paper, not separate read-outs.

• Prompt-injection detection. The log-probability of each answer at every layer is flattened into a vector, and an outlier detector is fitted on normal prompts only 2303.08112 §5.3. With Pythia 12B, the tuned lens with an isolation forest or LOF reaches AUROC 0.98 to 1.00 on five of nine tasks (author-reported) 2303.08112 tab. 1. The paper's table shows the logit lens below the tuned lens on those five tasks, except a tie on MNLI with LOF 2303.08112 tab. 1.
• Prediction depth. Spearman correlation between a prompt's prediction depth and the iteration at which the model learned it, on Pythia 12B. The tuned lens is higher on eight of eleven tasks; the logit lens is higher on MNLI, RTE and WinoGrande 2303.08112 tab. 2 2303.08112 §5.4.
• Causal basis extraction. A search for directions that the lens depends on most, with an erasure by mean ablation 2303.08112 §4.1 2303.08112 eq. 10 2303.08112 eq. 11. On Pythia 410M at layer 18 the paper reports Spearman rho = 0.89 between influence on the lens and influence on the model 2303.08112 §4.1 2303.08112 fig. 8.

When to use which

Implementation sketch

Language-agnostic pseudocode of the common core, with notes on where variants differ.

# Logit lens at layer l  [arxiv:2303.08112 §2] [arxiv:2303.08112 eq. 3]
def logit_lens(h_l):
    # Decode the hidden state with the final norm and unembedding  [§2]
    return softmax(unembed(final_norm(h_l)))

# Tuned lens at layer l  [arxiv:2303.08112 §3] [arxiv:2303.08112 eq. 8]
def tuned_lens(h_l, A_l, b_l):
    # Learned affine map, then the same final norm and unembedding  [eq. 8]
    return softmax(unembed(final_norm(A_l @ h_l + b_l)))

# Training loss for one batch of tokens  [arxiv:2303.08112 §3] [eq. 9]
def tuned_lens_loss(frozen_model, tokens, translators):
    # One forward pass that keeps every hidden state  [§3]
    target, hidden = frozen_model.forward_with_hidden(tokens)
    loss = 0.0
    # Each translator sees only its own layer; the last state is skipped  [§3]
    for l, h_l in enumerate(hidden[:-1]):
        # KL from the final-layer distribution to the lens output  [eq. 9]
        loss += kl(target, tuned_lens(h_l, *translators[l])).mean()
    return loss

Notes on where the variants differ:

• Initialisation. Translators start as the identity map 2303.08112 §3.
• Optimiser. The paper trains with SGD using Nesterov momentum, a linear learning-rate decay over 250 steps, gradient clipping at norm 1, weight decay 10^-3, and 2^18 tokens per optimiser step 2303.08112 §3.
• Learning rate. The paper states a base learning rate of 1.0 (0.25 when the final layer is kept) 2303.08112 §3. The released code passes 0.1 to a PyTorch SGD with momentum 0.9, and comments that PyTorch's scaling makes the effective step 1.0 ingredients.py. Copying 1.0 into torch.optim.SGD with momentum 0.9 gives a step about ten times too large. Momentum 0.9 is not stated in the paper ingredients.py.
• Weight decay target. The paper's A h + b form decays A toward zero. The released code uses a residual form that decays toward the identity 2303.08112 eq. 8 lenses.py. The two forms are equal at initialisation but regularise differently. Choose one and report it.
• Final layer. The paper includes the final layer for GPT-Neo and excludes it for GPT-2, GPT-NeoX-20B, OPT and Pythia 2303.08112 §3. The released training loop drops the last hidden state for every model train_loop.py. Reproducing the GPT-Neo lens needs a decision the paper does not spell out.
• Muon. The paper's later note reports that training with Muon converges much faster and reaches lower KL than SGD, but the authors did not redo the experiments with it 2303.08112 §3. The released code has no Muon implementation. Treat reported numbers as SGD-trained.

Contradictions

• Prompt-injection baseline. The text says the tuned lens beats the middle-layer baseline (SRM) on one task, ARC-Challenge 2303.08112 §5.3. The table shows a tie at 0.73 on ARC-Easy with LOF 2303.08112 tab. 1.
• Model editing. Reduced toxicity with the tuned lens on OPT-125m does not carry over to pythia-125m, pythia-350m, gpt-2-medium or gpt-neo-125m, where the paper finds no decrease above 2% 2303.08112 app. E.4 2303.08112 tab. 4.
• Earlier versions. The paper's section 3 states that lenses in earlier versions were severely undertrained 2303.08112 §3. Numbers from earlier versions should not be compared with the current ones.

Only one paper card covers this concept, so there is no cross-paper disagreement to report.

Papers

• Eliciting Latent Predictions from Transformers with the Tuned Lens: introduces the tuned lens, a per-layer affine read-out trained by KL distillation, and reports its use for prompt-injection detection, prediction depth and causal basis extraction.