PaperlayerView exact revision on arXiv

Source-linked AI summary

Eliciting Latent Predictions from Transformers with the Tuned Lens

Nora Belrose, Igor Ostrovsky, Lev McKinney, Zach Furman, Logan Smith, Danny Halawi, Stella Biderman, Jacob Steinhardt

arXiv:2303.08112v6cs.LG

TL;DR

The paper asks how transformer predictions are refined across layers and addresses the brittleness of directly decoding intermediate states. It trains layer-specific affine translators—the tuned lens—to produce vocabulary distributions from hidden states, finding more predictive, reliable, and unbiased trajectories than the logit lens, with applications including prompt-injection detection.

  • Problem

    Existing approaches provide limited or unreliable evidence about how pretrained transformers’ internal representations support iterative prediction, while the logit lens can be brittle across models and layers.

  • Method

    The tuned lens trains one affine translator for each layer and composes it with the pretrained unembedding to decode intermediate hidden states into vocabulary distributions.

  • Results

    The tuned lens produces substantially lower perplexity, lower bias, and outputs more representative of the final layer than the logit lens, while its prediction trajectories can detect prompt injection attacks often with near-perfect accuracy.

  • Takeaways & Limitations

    The tuned lens is a drop-in replacement for the logit lens that enables interpretable prediction trajectories in essentially any pretrained language model in use today.

  • Takeaways & Limitations

    The work focuses on language models because of space and time limitations, although the authors think the approach may apply to other modalities.

Abstract

from arXiv · show

We analyze transformers from the perspective of iterative inference, seeking to understand how model predictions are refined layer by layer. To do so, we train an affine probe for each block in a frozen pretrained model, making it possible to decode every hidden state into a distribution over the vocabulary. Our method, the tuned lens, is a refinement of the earlier "logit lens" technique, which yielded useful insights but is often brittle. We test our method on various autoregressive language models with up to 20B parameters, showing it to be more predictive, reliable and unbiased than the logit lens. With causal experiments, we show the tuned lens uses similar features to the model itself. We also find the trajectory of latent predictions can be used to detect malicious inputs with high accuracy. All code needed to reproduce our results can be found at https://github.com/AlignmentResearch/tuned-lens.

1. Introduction

The paper studies transformers as performing iterative inference, with each layer incrementally refining a latent next-token prediction. It introduces the tuned lens to decode these intermediate predictions more reliably and applies them to interpretability and prompt-injection detection.

  • Iterative inference: Transformer layers are treated as incremental updates to a latent prediction of the next token.Intermediate hidden states are decoded into vocabulary distributions, forming a layer-by-layer prediction trajectory.
  • Prediction trajectories: Tuned lens trajectories tend to converge smoothly toward the final output distribution, with successive layers achieving lower perplexity.This behavior supports using the trajectories to examine how predictions evolve during a forward pass.
  • Motivation: The logit lens can be unreliable because representations drift across layers, while other early-exit methods require changing the training process.The tuned lens is presented as a way to analyze pretrained models without modifying their training.
  • Method: The tuned lens trains one affine translator per layer to align intermediate representations with the final layer’s logits.Each translator is composed with the pretrained unembedding to map hidden states to vocabulary distributions.
  • Validation: The tuned lens has substantially lower perplexity than the logit lens and produces outputs more representative of the final layer distribution.Causal basis extraction further finds that features influential on tuned lens outputs are disproportionately influential on model outputs.
  • Applications: Prediction trajectories can detect prompt injection attacks often with near-perfect accuracy, and later-layer classification is associated with examples requiring many training steps.These are presented as applications of the tuned lens beyond qualitative inspection.

2. The Logit Lens

The logit lens directly decodes intermediate hidden states with the pretrained unembedding, but its reliability and interpretability vary across models and layers. The paper motivates the tuned lens by documenting representational drift, systematic bias, and high perplexity in logit-lens outputs.

  • Definition: The logit lens decodes hidden states into vocabulary space using the model’s pretrained unembedding matrix.It was previously used to inspect how intermediate predictions converge toward final answers.
  • Method: In a pre-LayerNorm transformer, the model is decomposed into layers up to an intermediate state and layers that map that state to logits.The logit lens approximates the latter computation by setting residual outputs to zero.
  • Reliability: The logit lens can fail to extract useful information from GPT-Neo, and extensions that retain the final transformer layer are only partially successful.Figure 1 provides an example of this limited recovery of meaningful results.
  • Reliability: The logit lens struggles with BLOOM and OPT 125M, where its top-1 prediction is often the input token rather than a plausible continuation in more than half the layers.The extracted information also varies across models and layers, complicating interpretation.
  • Bias: The logit lens systematically assigns excess probability mass to certain vocabulary items relative to the final output, making its trajectory a biased estimator.For GPT-Neo-2.7B, the bias is around 4 to 5 bits for most layers, compared with 0.0068 bits for Pythia 160M’s final distribution relative to Pythia 12B.
  • Comparison: Tuned lens predictions have significantly lower bias and much lower perplexity than logit-lens predictions across the evaluated settings.For BLOOM 560M, this remains true whether the final transformer layer is included in or excluded from the probe.
  • Bias: The paper measures bias using KL divergence between the marginal logit-lens and final-layer distributions.The unbiasedness criterion requires equal average probability for every vocabulary item across the dataset.

3. The Tuned Lens

The tuned lens addresses logit-lens brittleness by learning affine transformations that align intermediate representations with the final layer. It produces lower-bias, lower-perplexity predictions, transfers across nearby and fine-tuned layers, and supports iterative-inference analyses.

  • Method: The tuned lens replaces zero-residual decoding with a learned affine transformation for each layer.It learns a change-of-basis matrix and a constant offset, then decodes the transformed hidden state into vocabulary logits.
  • Method: It trains each translator by minimizing KL divergence between tuned-lens logits and the model’s final-layer logits.The final-layer distribution acts as a soft label, discouraging probes from learning information beyond the model’s own representations.
  • Training: Using Muon instead of SGD dramatically accelerates tuned-lens training and achieves much lower KL-divergence losses, indicating earlier lenses were severely undertrained.Muon-trained lens matrices also move farther from the logit lens, although the authors did not repeat all experiments with them.
  • Results: The tuned lens resolves logit-lens problems, with significantly lower bias and much lower perplexity across evaluated models.The evaluation uses validation data from several autoregressive language-model families, including Pythia and GPT-NeoX-20B.
  • Results: Tuned-lens translators usually transfer to nearby layers with modest perplexity increases, and transfer penalties correlate negatively with covariance similarity (Spearman ρ = −0.78).Transfer penalties are higher when probes trained on layers with outlier dimensions are evaluated on layers without them.
  • Results: Lenses trained on base models transfer well to fine-tuned versions without additional lens training, outperforming the logit lens and comparing well with model-specific lenses.The transferred lens uses the fine-tuned model’s unembedding while copying affine translators from the base-model lens.

4. Measuring Causal Fidelity

The paper evaluates whether tuned-lens features correspond to causally important model features and whether interventions produce aligned changes in intermediate and final predictions.

  • Causal basis extraction: Causal experiments test whether directions important to the tuned lens are also important to the final model output.The authors explicitly frame this as a desired property of probe-based interpretations.
  • Causal basis extraction: Causal basis extraction extends amnesic probing by finding multiple orthogonal directions that maximally degrade a function’s accuracy.The method iteratively searches for influential directions and orders them by influence.
  • Causal basis extraction: Spearman ρ = 0.89 between feature influence on the tuned lens and on the model at layer 18 of Pythia 410M.The results report no features influential in the tuned lens but not in the model, while the model is generally more causally sensitive.
  • Stimulus-response alignment: Stimulus-response alignment measures whether tuned-lens and final-layer distributions move similarly after the same hidden-state intervention.The comparison uses Aitchison differences and Aitchison similarity to control for absolute magnitudes.
  • Stimulus-response alignment: Stimuli align more with responses at later layers, and alignment is somewhat higher for tuned-lens than logit-lens principal subspaces.This pattern is reported for Pythia 160M across layers.
  • Stimulus-response alignment: The intervention uses resampling ablation on the subspace spanned by each layer’s top 10 causal basis vectors.These directions define the principal subspace used to compare logit- and tuned-lens responses.

5. Applications

The paper applies the tuned lens to interpret latent prediction trajectories, detect prompt injections, and estimate example difficulty in pretrained transformers. It reports strong detection performance and positive relationships between prediction depth and training difficulty, while noting that a baseline often performs better.

  • Prior applications: The tuned lens reproduces and extends prior logit-lens applications to downstream tasks and hidden knowledge elicitation.Its performance on taboo-word elicitation is mixed across layers, whereas downstream-task results recover prior early-layer robustness patterns.
  • Detecting Prompt Injections: Prompt-injection detection performance is nearly perfect when attacks drive accuracy well below random baseline and remains above chance for partially successful attacks.The table reports pooled AUROCs with 95% bootstrap confidence intervals over 10 random train-test splits.
  • Detecting Prompt Injections: The tuned lens anomaly detector achieves perfect or near-perfect AU-ROC on BoolQ, MNLI, QNLI, QQP, and SST-2.The detector is trained only on normal trajectories and evaluated on held-out normal and prompt-injected trajectories.
  • Detecting Prompt Injections: SRM consistently performs well: the tuned lens beats it only on ARC-Challenge, while SRM wins on MC TACO and SciQ.The tuned lens generally outperforms its logit-lens counterpart, but not this baseline across most tasks.
  • Measuring Example Difficulty: Prediction depth is the number of layers after which a prompt’s tuned-lens top-1 prediction stops changing.The measure estimates example difficulty without fine-tuning the model for early exiting.
  • Measuring Example Difficulty: A significant positive correlation with iteration learned appears on every investigated task, and tuned-lens correlations exceed logit-lens correlations in 8 out of 11 tasks.Iteration learned is the earliest training step at which a datapoint’s top-1 prediction becomes fixed.

6. Discussion

The tuned lens provides qualitative and quantitative insights into pretrained language models, including prompt-injection detection and causal feature analysis, while requiring per-layer translator training.

  • Contributions: The tuned lens is presented as a drop-in replacement for the logit lens that enables interpretable prediction trajectories in essentially any pretrained language model.The paper highlights prompt-injection detection as one initial application.
  • Contributions: Causal basis extraction identifies influential neural-network features and is proposed as potentially useful for interpretability research.The method is introduced as a separate contribution alongside the tuned lens.
  • Limitations: Training a translator layer for every network layer is the tuned lens’s main limitation relative to the out-of-the-box logit lens.The authors report that training a full probe set takes under an hour on a single 8×A40 node and provide checkpoints for commonly used models.
  • Limitations: Causal basis extraction is computationally intensive because it sequentially optimizes d_model causal basis vectors for each network layer.The authors suggest optimizing whole k-dimensional subspaces as a possible scalability improvement.
  • Scope: The work focuses on language models because of space and time limitations, although the authors think the approach may apply to other modalities.

A. Additional evaluation results

Additional evaluations compare tuned-lens variants and prediction perplexities across GPT-2, GPT-Neo, and OPT model families, with OPT 350M omitted because it uses post-LN architecture.

  • Lens comparisons: Figure 13 compares a Vicuna 13B-specific lens, a LLaMA 13B-transferred lens, and Vicuna 13B’s logit lens using KL divergence to the model’s final output.
  • Perplexity evaluation: Figure 14 reports perplexities of elicited predictions for OpenAI GPT-2, EleutherAI GPT-Neo, and Meta OPT models.
  • Perplexity evaluation: OPT 350M is omitted from the evaluation because it uses a post-LN architecture.

B. Qualitative Results

Qualitative trajectories show early confident predictions in some prompts, while logit-lens trajectories can exhibit pathological high probability on input tokens.

  • Tuned lens trajectories: Pythia 12B becomes highly confident at early layers after processing “It was the best of times,” suggesting some degree of memorization.
  • Tuned lens trajectories: Figure 16 presents a tuned-lens prediction trajectory for Pythia 12B prompted with the abstract of Brown et al. (2020).
  • Logit lens pathologies: The BLOOM 560M logit lens assigns very high probability to the input token at many layers and positions, complicating output interpretation.
  • Logit lens pathologies: The OPT logit-lens evaluation includes the facebook/opt-125m model.
  • Logit lens pathologies: The OPT 125M logit lens exhibits a similar pathology to BLOOM 560M.

C. Transformers Perform Iterative Inference

Residual connections support an iterative-inference view in which layers update representations toward lower loss, and experiments with Pythia 6.9B find consistently negative residual-gradient alignment.

  • Theory: Skip connections encourage neural networks to perform iterative inference by updating hidden states in directions of decreasing loss.The theoretical argument is presented as applying to transformers and other residual networks.
  • Theory: A residual block updates a representation h_i through its residual transformation.
  • Theory: Taylor expansion of the final loss around an intermediate representation provides the first-order basis for analyzing residual updates.
  • Theory: The analysis labels the relationship between the residual and gradient as gradient-residual alignment.
  • Theory: To first order, minimizing the inner product between the residual and loss gradient can be achieved by aligning the residual with the negative gradient.
  • Empirical evaluation: Empirical measurement uses cosine similarity between each residual and the corresponding loss gradient.
  • Empirical evaluation: Residual-gradient cosine similarity is consistently negative, especially in the final stage of a ResNet image classifier.
  • Empirical evaluation: For Pythia 6.9B, residual-gradient cosine similarity is negative in every layer at least 95% of the time.Observed magnitudes never exceed 0.05 but are much larger than expected for random vectors in the high-dimensional space.

C.1. Zero-shot robustness to layer deletion

The experiment tests whether transformer performance remains robust when layers are deleted, extending layer-deletion findings from ResNets to transformers. In the tested transformer, the first layer is uniquely important while later-layer deletions have little effect.

  • The experiment applies stochastic depth, also called LayerDrop, as a regularization technique that randomly drops layers during training.
  • ResNets remain robust to layer deletion without stochastic depth, whereas CNNs without skip connections do not.
  • Only the first layer is crucial for performance; deleting every other layer causes a nearly imperceptible perplexity increase.

D.1. Ablation techniques

The paper contrasts naive projection-based ablation with mean-replacement ablation to avoid out-of-distribution activations. This adjustment is intended to make causal interventions more interpretable while preserving the model’s typical activation statistics.

  • Projecting onto a subspace’s orthogonal complement removes information along that subspace by enforcing zero inner products with its directions.
  • Naive projection can increase loss by removing implicit bias terms, even when the projected component carries little useful input information.
  • Mean-replacement ablation substitutes a direction’s dataset-wide mean value instead of zeroing that direction.
  • The tuned lens is evaluated in static analysis and model editing, where it does not decrease static-analysis performance and outperforms the logit lens on OPT-125m toxicity reduction.

E.1. Static Analysis

The paper evaluates whether intermediate transformer parameters yield interpretable token patterns under tuned- and logit-lens projections. The tuned lens generally improves interpretability comparisons, though benefits weaken in larger models and toxicity-editing results are limited across several models.

  • Parameter interpretability is assessed by projecting extracted hidden-state vectors into token space and selecting the top k associated tokens.
  • The tuned lens replaces direct unembedding with a layer-specific translator that projects an intermediate hidden state toward the final hidden-state basis.
  • Tuned- and logit-lens projections make parameters more interpretable than random matrices on Pythia 125M, with the tuned lens showing a larger real-versus-random score gap.
  • Most singular vectors are no more interpretable than random, while a minority form a long right tail, motivating comparison against randomly generated matrices.
  • A BPEmb-based metric averages pairwise token-embedding cosine similarities to approximate the monosemanticity of each token set.
Loading 2303.08112v6…