Transformer layers often undo each other to improve language models
LLM Layers Immediately Correct Each Other
Computation and LanguageMachine Learning
Summary
Language models like transformers process words through multiple layers, each adding information to understand text better. The authors found that instead of just building on previous layers, many layers actually correct or undo parts of what the previous layers did. This back-and-forth helps the model avoid mistakes and fine-tune its understanding. They explain this as layers proposing ideas and then later layers rejecting the ones that don't fit. This discovery changes how we think about how these models learn features inside them.
transformerlanguage modellayerresidual streamfeaturesparse autoencoderpretrainingJacobianmodel interpretabilitycontextual tokens
Authors
Arjun Patrawala, Jiahai Feng, Erik Jones, Jacob Steinhardt
Abstract
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.