Source-linked AI summary
LLM Layers Immediately Correct Each Other
Arjun Patrawala, Jiahai Feng, Erik Jones, Jacob Steinhardt
TL;DR
Interpretability methods often treat layer outputs as accumulating persistent features, but evidence for systematic interactions between adjacent layers is limited. This paper characterizes the Transformer Layer Correction Mechanism using cross-model observations, causal interventions, and layer-Jacobian analysis. TLCM appears in 5 of 7 open-weight model families and supports a propose-and-reject interpretation with implications for feature specificity and steering.
Problem
Interpretability methods commonly view transformer layers as additively building persistent residual-stream features, leaving adjacent-layer correction insufficiently characterized.
Method
The paper measures adjacent-layer contribution similarities, uses causal interventions to test dependence, and applies the layer Jacobian to identify corrected and reinforced subspaces.
Results
TLCM occurs in 5 of 7 open-weight model families, emerges during pretraining, activates especially on context-dependent tokens, and adaptively depends on prior-layer output.
Takeaways & Limitations
The propose-and-reject interpretation suggests that residual streams contain both persistent features and transient proposals, helping explain low SAE feature specificity and extreme steering amplification.
Takeaways & Limitations
The paper does not explicitly test the propose-and-reject hypothesis.
Abstract
from arXiv · showhide
Recent methods in language model interpretability employ techniques such as sparse autoencoders to decompose residual stream contributions into linear, semantically meaningful features. Such methods are commonly interpreted as identifying features that persist in the residual stream and that subsequent layers build upon. We challenge this view by identifying the Transformer Layer Correction Mechanism (TLCM), wherein adjacent transformer layers systematically counteract portions of each other's contributions. TLCM appears in 5 out of 7 major open-source model families and activates across nearly all tokens in diverse texts. We show that TLCM emerges during pretraining, operates most strongly on contextually dependent tokens, and adaptively calibrates its correction strength based on the preceding layer's output. Using the layer Jacobian, we further show that TLCM selectively corrects specific subspaces while reinforcing others, which we interpret through a ``propose-and-reject'' framework in which layers propose candidate features and subsequent layers selectively remove inappropriate ones. This dynamic suggests that the residual stream at any layer contains transient proposals alongside persistent features, helping explain why SAE feature descriptions often have low specificity, why effective model steering requires extreme feature amplification, and why transcoders hold a theoretical advantage over SAEs.
1 Introduction
The paper challenges the view that transformer layers simply accumulate representations by identifying systematic adjacent-layer corrections. TLCM appears across many model families and motivates a propose-and-reject interpretation of the residual stream.
- 1 Introduction: Recent interpretability methods often assume that transformer layers additively enrich residual-stream representations and build persistent features.This view motivates probes, logit lens analyses, SAEs, and cross-layer transcoders.
- 1 Introduction: TLCM occurs in 5 of 7 open-weight LLM families, with adjacent layers systematically reversing portions of one another’s contributions.The identified families are Llama 3, OLMo, Mistral, Gemma, and Qwen2.
- 1 Introduction: Adjacent-layer cosine similarity can reach −0.4, indicating that some layer output reverses preceding content rather than adding new content.Single-layer features may therefore be corrected by the next layer or function as corrections themselves.
- 1 Introduction: TLCM emerges during pretraining, activates most frequently on context-dependent tokens, and reflects contributions from both attention and MLP layers.Reported examples include numbers, dates, and punctuation.
- 1 Introduction: Causal interventions show that layer i + 1 adaptively adjusts correction strength according to layer i’s contribution, while Jacobian analysis identifies selectively corrected and reinforced subspaces.A subspace is corrected 1-for-1, a smaller subspace is reinforced, and non-TLCM interactions show neither pattern.
- 1 Introduction: These findings support a propose-and-reject view in which layers propose candidate features and later layers selectively remove inappropriate ones.The framework connects TLCM to SAE specificity, model steering amplification, and the possible advantage of transcoders over SAEs.
2 Related Work
Related work frames transformer representations as progressively refined and examines both feature decomposition and correction-like behavior. TLCM extends these discussions to pervasive corrections during ordinary forward passes.
- 2 Related Work: Feature-based interpretability extracts semantic features from neural activations using supervised probes and unsupervised methods such as SAEs.These methods study the role of features in layer outputs and residual-stream representations.
- 2 Related Work: Prior self-repair work studies later components compensating for earlier ablations or layer swaps, including mechanisms such as copy suppression.TLCM differs by occurring repeatedly on nearly all tokens during an ordinary forward pass.
- 2 Related Work: The Iterative Inference Hypothesis and enrichment studies describe successive layers as progressively refining latent representations in the residual stream.Examples include function vectors placed in the residual stream to induce behavior.
3 Background
The paper defines transformer contributions to the residual stream and introduces the similarity-matrix view used to identify layer corrections. It also situates SAEs as a feature-decomposition method and specifies the studied models.
- 3 Background: LLMs embed tokens in a d_m-dimensional residual space, process them through n transformer layers, and then unembed the resulting representations.Positional embeddings are added before the transformer stack.
- 3 Background: Figure 1 measures average cosine similarity between residual-stream contributions from layer pairs, with negative values indicating opposing vectors and positive values indicating reinforcement.Adjacent-layer blue bands reveal systematic reversals in some models, while their absence indicates no consistent reversals.
- 3 Background: Each transformer layer contains attention and MLP sublayers that contribute vectors to the residual stream.The marginal contribution is c_i = Attn_i(x_i) + MLP_i(x_i + Attn_i(x_i)), with x_i+1 = x_i + c_i.
- 3 Background: Sparse autoencoders decompose residual-stream contributions into a small number of orthogonal feature vectors with activations plus an error term.The number of features k is usually less than 500, and the features can be semantically labeled and manipulated.
- 3 Background: The study evaluates TLCM across 7 open-weight model families, primarily reporting results on Llama 3.1 8B.The families include Llama 3, OLMo, Mistral, Gemma, Qwen2, GPT-2, and Phi.
4 Transformer Layer Correction Mechanism
The paper identifies TLCM as a recurring partial reversal between adjacent transformer-layer contributions, challenging the view that layers mainly enrich the residual stream additively. TLCM emerges during pretraining, varies with context and token type, and involves coordinated attention and MLP behavior.
- TLCM defines adjacent-layer contributions that systematically oppose one another, with negative cosine similarity averaging approximately −0.2 across a large corpus.
- TLCM appears in 5 of 7 open-weight model families, typically with adjacent-layer cosine similarities from −0.3 to −0.15 and as low as −0.4.It persists across text types, model scales, and instruction-tuned or conversational models, but is absent in GPT-2 and Microsoft Phi models.
- The cosine-similarity distribution is bimodal, with a negative mode near −0.2 and a valley around −0.1 used to define TLCM activations.The distinct negative mode indicates a recurring interaction pattern rather than rare outliers.
- Independent-output null models place typical TLCM instances roughly 3 to 6 standard deviations from expectation, although these comparisons do not establish causality.The standard-normal null gives a 6-sigma threshold at d = 4096, while the empirical-layer null gives roughly 3 sigma.
- TLCM emerges during pretraining rather than appearing at initialization or through later supervised or reinforcement fine-tuning.In OLMo, negative adjacent-layer structure is absent at 5% of training, begins around 21%, and strengthens thereafter.
- MLP sublayers produce anti-correlated contributions with preceding attention and MLP sublayers, whereas attention sublayers correlate positively with other attention sublayers.These patterns suggest that MLPs execute much of the correction, while the mechanism remains a transformer-layer-level phenomenon.
5 TLCM Adaptivity
TLCM is adaptively driven by the preceding layer and selectively corrects particular contribution subspaces rather than uniformly attenuating them. These findings motivate a propose-and-reject view in which layers propose features and subsequent layers remove inappropriate ones.
- 5.1 TLCM is Adaptive: Causal interventions show that layer i + 1 directly adjusts its correction strength in response to layer i’s contribution.Increasing the perturbation strengthens correction, while decreasing it weakens correction.
- 5.1 TLCM is Adaptive: At α = −1, adjacent-layer similarity becomes positive, indicating compensatory boosting that is strongest in earlier layers.Extreme perturbations place the model in counterfactual states.
- 5.2 Isolating the Correction Subspace: The layer Jacobian decomposes input directions into corrected, reinforced, or approximately untouched subspaces.Negative eigenvalues indicate correction, positive eigenvalues reinforcement, and near-zero eigenvalues little change.
- 5.2 Isolating the Correction Subspace: A TLCM eigenvalue mode near λ = −1 indicates unit-sized reversal, while a mode near λ = 0.5 indicates reinforcement.These results come from aggregated Jacobian analyses of TLCM and non-TLCM pairs.
- 5.2 Isolating the Correction Subspace: About 3000 of 4096 Jacobian eigenvalues are negative, while TLCM’s bimodal distribution shows selective correction rather than uniform attenuation.The analysis uses the top fifty directions of each prior-layer contribution, accounting for 15+% of its variance.
- 5.3 Propose and Reject Hypothesis: The propose-and-reject hypothesis describes a layer proposing candidate features followed by a context-aware subsequent layer removing irrelevant ones.The framework is consistent with attention and MLP participation, context-position effects, and selective correction, but is not explicitly tested.
6 Discussion
TLCM reframes residual-stream features as a mixture of persistent features and transient proposals that later layers may reverse. This perspective connects TLCM to SAE specificity, steering amplification, and the interpretability advantages of transcoders, while leaving several questions open.
- Feature descriptions: Single-layer SAEs cannot distinguish persistent features from soon-to-be-reversed proposals, helping explain low-specificity feature descriptions.Over 50% of activated SAE features in a cited study were labeled “Irrelevant” or “Only vaguely related.”
- Model steering: Steering interventions may need to exceed TLCM’s correction capacity, which begins diminishing when the prior layer is amplified beyond 2×.The passage relates this to reports of steering requiring amplification up to 10× a feature’s maximum observed value.
- Transcoders: Transcoders can in principle represent corrections as input-conditional mechanisms that distinguish corrected features from newly contributed features.This is offered as an interpretation of their reported reconstruction and interpretability advantage over SAEs.
- Open questions: Whether TLCM is an efficient mechanism or a consequence of architectural constraints remains open, as does its presence in other sequence-model architectures.The passage specifically names SSMs and transformer variants such as gated attention, NoPE, and MOE.
- Open questions: The propose-and-reject hypothesis remains a conceptual framework without formalization.The paper suggests circuit tracing and automated feature labeling as possible routes toward semantic explanations of correction.
- Implications: TLCM may serve as a constraint for researchers developing mechanistic accounts of transformer computation.
A LayerNorm Blindness to explain TLCM
RMSNorm can make attention and MLP sublayers blind to residual-stream magnitude, offering one possible source of over-contribution and subsequent correction. However, LayerNorm blindness alone does not explain TLCM, which persists under alternative normalization designs.
- RMSNorm mechanism: RMSNorm makes attention and MLP sublayers blind to the residual stream’s norm.This blindness can cause under-contribution when the residual norm is high and may encourage over-contribution followed by correction.
- Limits of the explanation: LayerNorm blindness alone does not fully explain TLCM because the mechanism does not entirely reverse the prior contribution.
- Alternative normalization designs: TLCM remains robust in Gemma 2 despite its hybrid pre-LayerNorm and post-LayerNorm architecture.Post-LayerNorm should in principle make residual-stream norm more predictable.
- Alternative normalization designs: The correction mechanism persists strongly in OLMo 2 despite replacing pre-LayerNorm with post-LayerNorm and QK norm.
B Jacobian Eigendecomposition Locates Corrected Subspaces
The Jacobian eigendecomposition identifies which perturbation directions are corrected, reinforced, or largely untouched by the following transformer layer.
- Corrected directions: Negative-eigenvalue eigenvectors correspond to corrected subspaces, with more negative eigenvalues indicating stronger correction.Correction occurs when the perturbation has sufficient projection onto directions with negative eigenvalues.
- Reinforced directions: Positive-eigenvalue eigenvectors are reinforced in proportion to their eigenvalues.
- Untouched directions: Near-zero-eigenvalue eigenvectors are mostly untouched by the layer.
- Jacobian decomposition: The Jacobian is symmetrized and decomposed as J = QVQ^⊤, with eigenvectors defining directions and eigenvalues quantifying their effects.The projected perturbation Q^⊤∆ weights each eigendirection in the quadratic correction condition.
C.2 Experimental Mean and Standard Deviation
Layer contributions are concentrated in a low-dimensional, anisotropic subspace, and the empirical TLCM cutoff is rare under cross-token variation.
- Contribution anisotropy: 500-800 of Llama’s 4096 residual dimensions capture the main concentration of a layer’s contribution vectors.This concentration is estimated from singular-value decompositions of contribution vectors across 4096 WikiText tokens.
- Empirical distribution: −0.00375 mean and 0.03267 standard deviation characterize adjacent-layer cosine similarity across different tokens.
- Statistical cutoff: The −0.1 TLCM cutoff is approximately a 3σ event under this cross-token distribution.The distributions are plotted by layer in Figure A3.
D Details on Jacobian Experiments
The Jacobian analysis relies on a locally accurate approximation and examines anisotropic contribution directions and cross-token cosine-similarity distributions.
- Contribution structure: In TLCM-active layers 4-20, the first 500-800 principal components explain 80% of contribution variance.The blue curve reports explained variance, while the orange curve represents the isotropic expectation.
- Jacobian sanity check: For |α| < 0.1, the Jacobian approximation has consistently around 5% error across 46 randomly selected TLCM Jacobians.The perturbations use ∆ = αa_t over α values from −0.5 to 0.5.
- Analysis basis: The Jacobian approximation is therefore treated as sufficiently accurate for decomposing directions in the subsequent analysis.
E Extended Details on TLCM Existence
Extended experiments examine TLCM across models and characterize how attention and MLP components contribute to its corrections and adaptive behavior.
- Model coverage: TLCM’s existence is evaluated across many HuggingFace models using the same technique as the main analysis.
- Component interactions: MLPs correct prior attention and MLP contributions, whereas attention corrects only prior MLP contributions.The authors suggest a shared low-dimensional communication subspace may contribute to these interactions.
- Adaptive correction: TLCM’s adaptive correction is plotted across different layers using the main-text procedure.
F.2 TLCM Correction Capacity Diminishes at High α
TLCM correction strength increases with the prior layer’s contribution but begins to diminish when that contribution is scaled to extreme values.
- At α > 2, TLCM’s correction capacity starts to diminish as the prior layer is scaled to extremely large values.The paper suggests this may reflect saturation of the correction mechanism.
- The experiments visualize TLCM matrices across multiple models while clamping values to the range −0.2 to 0.2 and zeroing diagonal entries.The displayed model examples include Qwen2.5-0.5B-Instruct and OLMo-2-1124-7B-Instruct.
G OLMo Training Checkpoint Experiment Details
The OLMo checkpoint experiments track TLCM across training stages and show that the mechanism emerges progressively during pretraining. The analysis uses checkpoint-specific matrices and a corpus of roughly 500 PyTorch-instructional tokens, with corpus size becoming negligible beyond a sufficient token count.
- Experiment setup: The OLMo 7B checkpoint analysis computes M on a corpus of about 500 tokens of PyTorch instructional content.The corpus size has little effect on cosine-similarity matrices beyond a sufficiently high token count.
- OLMo training checkpoints: TLCM emerges progressively during OLMo 7B pretraining, first appearing around step 135,500 after 0.5T tokens.The rightmost plot represents the fully pretrained model.
- Visualization: The checkpoint visualizations clamp M to −0.2 through 0.2 and zero diagonal entries for clarity.The correction-strength plots separately sample random TLCM curves across layers.
- OLMo training checkpoints: TLCM also emerges progressively during pretraining in OLMo 2 1B and OLMo 2 7B.The figures compare checkpoints at different training steps and token counts.