---
kind: concept
slug: logit-lens
title: Logit lens
summary: >-
  Reads a transformer's intermediate hidden states as next-token guesses by applying the final norm and unembedding. The tuned lens adds one trained affine map per layer to match the final output. On Pythia 12B, a tuned-lens detector reaches AUROC 0.98 to 1.00 on five of nine prompt-injection tasks (author-reported); the plain logit lens shows higher bias on several models.
tags: [interpretability, linear-probes, transformers]
papers: ["2303.08112"]
related: [mechanistic-interpretability, linear-probing, prompt-injection, automated-interpretability]
status: draft
updated_at: 2026-11-09
---

## Overview

A transformer's hidden states can be read as successive refinements of a guess
for the next token [arxiv:2303.08112 §1] [arxiv: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 [arxiv: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 [arxiv:2303.08112 §2] [arxiv: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 [arxiv:2303.08112 §2] [arxiv: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
[arxiv:2303.08112 §2] [arxiv:2303.08112 fig. 3]. Card:
[Eliciting Latent Predictions from Transformers with the Tuned Lens](/papers/2303.08112).

### Logit lens with the last update kept

A variant keeps the final layer's residual update instead of dropping it
[arxiv:2303.08112 eq. 4] [arxiv:2303.08112 §2]. The paper reports that it
works better for some models and not for others [arxiv:2303.08112 fig. 1]
[arxiv: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 [arxiv:2303.08112 §3] [arxiv: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 [arxiv: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 [arxiv:2303.08112 §3]
[arxiv: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 [arxiv:2303.08112 §3].
Reported results: lower perplexity than the logit lens across the models
tested, with or without the final layer in the probe [arxiv:2303.08112 §3]
[arxiv:2303.08112 fig. 4] [arxiv: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 [arxiv:2303.08112 §3]
[arxiv: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 [arxiv: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) [arxiv: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 [arxiv: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 [arxiv:2303.08112 tab. 2] [arxiv:2303.08112 §5.4].
- **Causal basis extraction.** A search for directions that the lens depends
  on most, with an erasure by mean ablation [arxiv:2303.08112 §4.1]
  [arxiv:2303.08112 eq. 10] [arxiv: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 [arxiv:2303.08112 §4.1] [arxiv:2303.08112 fig. 8].

## When to use which

| Situation | Use | Why |
|---|---|---|
| No training budget, or the read-out must be exact and label-free | Logit lens | Needs no parameters. Check the bias for your model first: the paper reports it can be large for BLOOM, OPT 125M and GPT-Neo-2.7B [arxiv:2303.08112 §2] |
| Need comparable per-layer guesses across models, and a small probe is affordable | Tuned lens | Lower perplexity than the logit lens across the models tested [arxiv:2303.08112 §3] [arxiv:2303.08112 fig. 4]. Training reported under an hour on an 8 x A40 node (author-reported) [arxiv:2303.08112 §6] |
| Detecting tampered prompts from the layer-wise answer trajectory | Tuned lens with an outlier detector | Five of nine tasks at AUROC 0.98 to 1.00 on Pythia 12B (author-reported) [arxiv:2303.08112 tab. 1]. Fitted on normal prompts only, against fixed injected strings [arxiv:2303.08112 §5.3] |
| Only the final model is available, or you cannot retrain after a fine-tune | Tuned lens from the base model, checked on the fine-tune | Transfer loss reported as small, but the paper's check covers only the datasets it names [arxiv:2303.08112 app. A] |
| Mapping a layer to its output for a model not in the paper | Either, then measure bias | The paper reports results on the models it tested; no general guarantee is given [arxiv:2303.08112 §3] |

## Implementation sketch

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

```pseudocode
# 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 [arxiv: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 [arxiv:2303.08112 §3].
- **Learning rate.** The paper states a base learning rate of 1.0 (0.25 when
  the final layer is kept) [arxiv: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 [code: tuned_lens/scripts/ingredients.py#L247-L250].
  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
  [code: tuned_lens/scripts/ingredients.py#L210].
- **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
  [arxiv:2303.08112 eq. 8] [code: tuned_lens/nn/lenses.py#L306-L311]. 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 [arxiv:2303.08112 §3]. The
  released training loop drops the last hidden state for every model
  [code: tuned_lens/scripts/train_loop.py#L403]. 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 [arxiv: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 [arxiv:2303.08112 §5.3].
  The table shows a tie at 0.73 on ARC-Easy with LOF [arxiv: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% [arxiv:2303.08112 app. E.4]
  [arxiv:2303.08112 tab. 4].
- **Earlier versions.** The paper's section 3 states that lenses in earlier
  versions were severely undertrained [arxiv: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](/papers/2303.08112): 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.