Eliciting Latent Predictions from Transformers with the Tuned Lens

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

516 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Model Architectures, Safety, Robustness & Interpretability

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

In Safety, Robustness & Interpretability

Citations in the corpus

Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling (2023, methodology): “We use untied embedding / unembedding matrices, as prior work has suggested that this makes interpretability research easier (Belrose et al., 2023).”
The MiniPile Challenge for Data-Efficient Language Models (2023, background): “structurally different architectures [16, 63, 32], pre-training schemes [50, 41], optimizers schemes [3, 26, 24], differential privacy [2, 5], mechanistic interpretability [15, 46, 6], etc.”
Not What You've Signed Up For: Compromising Real-World LLM-Integrated Applications with Indirect Prompt Injection (2023, unknown): “A final promising solution is to rely on interpretability-based solutions that perform outlier detection of prediction trajectories [38].”
LLM Censorship: A Machine Learning Challenge or a Computer Security Problem? (2023, methodology): “For example, in this paper we only consider censorship based on inputs and outputs of a model and do not consider methods that rely on internal representations of LLMs [4, 7, 27].”
Do Llamas Work in English? On the Latent Language of Multilingual Transformers (2024, methodology): “(Plotting tool: Belrose et al. (2023).) so well from their mainly English training data to other languages?”
Emergent Linear Representations in World Models of Self-Supervised Sequence Models (2023, background): “Figure 6 visualizes OthelloGPT’s “iterative infer-ence” (Jastrzebski et al., 2018; Belrose et al., 2023; Veit et al., 2016; nostalgebraist, 2020), or iterative reﬁnement of features.”
Emergent and Predictable Memorization in Large Language Models (2023, methodology): “Specifically, these small/partial runs provide a cheap method to inform training behavior for a given corpus, rather than training an entire model from scratch.”
Mechanistic Interpretability for AI Safety - A Review (2024, background): “Path patching extends the activation patching approach to multiple edges in the computational graph (Wang et al., 2023; Goldowsky-Dill et al., 2023), allowing for a more fine-grained analysis of component interactions.”
Weight Tying Biases Token Embeddings Towards the Output Space (2026, unknown): “The tuned lens (Belrose et al., 2023) addresses this limitation by learning an affine transformation for each layer: where A ℓ ∈ R d d is a learned linear map and b ℓ ∈ R d is a bias term.”