Large Language Models (LLMs) are trained with next-token prediction, implemented in autoregressive Transformers via causal masking for parallelism. This creates a subtle misalignment: residual connections tie activations to the current token, while supervision targets the next token, potentially propagating mismatched information if the current token is not the most informative for prediction. In this work, we empirically localize this input-output alignment shift in pretrained LLMs, using decoding trajectories over tied embedding spaces and similarity-based metrics. Our experiments reveal that the hidden token representations switch from input alignment to output alignment deep within the network. Motivated by this observation, we propose a lightweight residual-path mitigation based on residual attenuation, implemented either as a fixed-layer intervention or as a learnable gating mechanism. Experiments on multiple benchmarks show that these strategies alleviate the representation misalignment and yield improvements, providing an efficient and general architectural enhancement for autoregressive Transformers.
Figures & tables
Fig. 1: Schematic illustration of the token misalignment between input and output in Transformer architectures (here with a single layer). Token ti in the input is directly connected to token ti+1 in the output through the residual connections.
Layer
token #107
token #108
token #109
token #110
token #111
token #112
token #113
Input
musical
chord
into
the
intensities
of
its
0 - 6
musical
chord
into , < bos >
the
intensities , myself
of , < bos >
its
7 - 9
musical
chord
AnchorStyles, These
the
itself
lenker
its
10
musical
chords
itself
the
itself
ReusableCell
####
11
musical
itself
AddTagHelper
various
itself
BoxDecoration
####
12
musical
itself
lenker
WithMany
itself
lenker
WebVitals
TABLE I: Decoded hidden states of Gemma-2-2B on a Wikipedia article about the Fourier transform. Layers are grouped by similar decodings: early layers (0–6) align with the input (green), while late layers (21–26) align with the one-token-shifted targets (red). For instance, token #107 consistently decodes to “musical” across layers 0–6, whereas token #109 alternates between “into” and <bos> within the same range.
Fig. 2: Average top-5 match rate between the decoded hidden states and either the input sequence (blue) or the shifted input sequence (red) as a function of layer depth in the Gemma-2-2B model. Results are averaged over 1,000 sequences from the Wikitext dataset. The shift from input to output indexing occurs late in the architecture.
Fig. 3: Continuous similarity measures between hidden states and their corresponding input and shifted input tokens, reported for Gemma-2-2B [ 15 ] , Llama-3.2-3B [ 5 ] and Mistral-7B-v0.3 [ jiang2023mistral7b ] . Metrics include cosine similarity to input and output token embeddings, and normalized projection along the input–output axis (0: input token, 1: output token). Results confirm the latent transition from input-based to output-based representations across different architectures.
Fig. 4: Impact of cutting the residual path at different layers in GPT2-0.1B, trained on 10B tokens from Fineweb [ penedo2024the ] on the validation loss. The baseline loss is represented with a dotted line.
Fig. 5: Evolution of the learned probability distribution of cutting the residual path at a given layer, during training. At the beginning, all 12 layers have the same probability of cut, then the 11th layer attracts all the mass of the probability.
benchmark
baseline
fixed-cut
gating
Wikitext
28.46
28.62
28.75
LAMBADA
32.99
33.38
33.52
OpenWebText
34.74
34.84
35.50
Fineweb-Edu
38.87
38.98
39.19
TABLE II: Comparison of the baseline, fixed attenuation, and learnable gating across multiple benchmarks: Wikitext [ 8 ] , LAMBADA [ 11 ] , OpenWebText [ 4 ] , and Fineweb-Edu dataset [ penedo2024the ] . Bold scores are the best ones.
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.
In contrast to RNNs, which compress their history into a single hidden state, Transformers can attend to all past tokens directly. However, standard Transformers rely solely on the hidden state from the previous layer to represent the entire context. We show that this design creates pressure toward representation collapse and can degrade performance. To address this issue, we introduce Layer-Integrated Memory (LIMe), a lightweight extension that leverages existing key-value buffers and learns per-head, per-layer routing weights to integrate representations from previous layers. Across language modeling, synthetic reasoning, and deep architectures, LIMe improves perplexity per FLOP in the studied regimes and yields strong gains on synthetic tasks while preserving higher value-vector entropy and token separability. Finally, learned routing weights reveal systematic reuse of local and long-distance features, showing how LIMe enriches attention-time memory without increasing hidden-state size. Code is available at https://github.com/corl-team/lime.
We study Latent Recurrent Transformer (LRT), a lightweight augmentation of autoregressive transformers that reuses a high-level source-layer hidden state from the previous token as recurrent memory for the next token. Because this state is already computed during ordinary decoding, LRT introduces a cross-token, cross-layer latent pathway while preserving the standard attention mechanism, KV-cache interface, and one model forward per generated token. To pretrain this recurrence without sequentially unrolling the full sequence, we introduce interleaved parallel training: one full-sequence initialization forward constructs a shared buffer, followed by sequential refinement of disjoint position subsets with parallel computation within each subset. This provides every token with recurrent-memory-aware supervision at approximately 2x ideal token compute. Across 1.3B- and 2.1B-parameter nanochat-style backbones and a wide range of training budgets, LRT improves both BPB and CORE under matched effective compute. Additionally, LRT outperforms two-forward PonderLM-2 and matches a three-loop Transformer in BPB, while retaining one-forward-per-token decoding with 9% latency overhead over the standard Transformer.
Zeyi Huang, Xuehai He, LiLiang Ren +8
Microsoft · University of Washington · University of Wisconsin-Madison