You Do Not Fully Utilize Transformer's Representation Capacity

The single stream is a bottleneck

A transformer never rereads its past layers. At layer $\ell$ the attention block sees exactly one object, the residual stream $\mathbf{X}_{\ell-1}$, and every query, key, and value is a projection of it. Whatever layers $1$ through $\ell-2$ computed survives only as far as it fits inside that single vector per position.

Forcing everything through one vector has a measurable failure mode: representation collapse, where distinct tokens become indistinguishable in deeper layers. Voita and colleagues showed top layers losing fine-grained token identity, and Barbero and colleagues proved a decoder can drive the final representations of different inputs arbitrarily close together. Long sequences make it worse, because subtle distinctions get squeezed out by finite state capacity and floating-point precision.

Our fix, Layer-Integrated Memory, or LIMe, removes the bottleneck without touching the transformer's shape. Every layer already computes per-head keys and values, and those tensors already sit in buffers during training and inference. We simply let every attention head read all of them.

A router over memory the model already has

At layer $\ell$, head $h$ no longer attends over the current layer's keys and values alone. A trainable router $R^{(\ell)}\in\mathbb{R}^{\ell\times H\times H}$ assigns a weight to every buffered head $h'$ at every earlier layer $\ell'$, and attention runs over the resulting mixtures

Queries stay untouched, masking stays causal, and FlashAttention still applies. We initialize the current layer's slice of the router to the identity, so training starts from an exact standard transformer and learns where to deviate. Giving the router its own learning rate of $10^{-2}$, with no weight decay, speeds up circuit formation.

Because the buffers are already allocated, the overhead is nearly unmeasurable. With grouped-query attention, LIMe adds 0.00 percent parameters, 0.08 percent forward FLOPs, 1.16 percent step time, and 0.00 percent peak training memory. Hyper-Connections, the strongest alternative that widens the residual stream instead of rereading it, costs over 23 percent extra step time in the same setup.

Faster convergence at one billion parameters

We pretrained 1B-parameter LLaMA-style models from scratch on roughly 50B tokens of FineWeb Edu. Measured in training FLOPs, LIMe reaches the baseline's loss 15.3 percent sooner, 8.9 percent with grouped-query attention, and finishes with 1.15 percent lower perplexity, 0.91 percent under GQA. Same data, same optimizer, same budget; the only difference is that attention may reread earlier layers.

The gains concentrate where recalling earlier features matters most. Key-Value Maps climbs from 45.9 to 55.6 and MultiRC from 43.2 to 56.2, while saturated pattern-completion tasks barely move. Across all eight tasks the average rises from 48.2 for LLaMA and 49.9 for Hyper-Connections to 51.7 for LIMe.

Watching collapse recede

To observe collapse directly, we borrowed a probe from Voita and colleagues: 1,668 occurrences each of is, are, was, and were, drawn from FineWeb Edu, with a four-way linear classifier trained on each layer's representations under five-fold cross-validation. If deep layers keep these near-synonyms distinct, the probe succeeds; if representations collapse, it cannot.

The division of labor reverses exactly as designed. LIMe's value vectors stay separable with probe accuracy near 1.0, and their matrix entropy stays higher than LLaMA's layer after layer. The residual stream itself becomes less separable than the baseline's, which is the point: the stream no longer has to carry every lexical nuance forward, because the buffers already do.

Where representation quality decides the answer

Language-model perplexity averages over mostly easy predictions, so we also measured tasks that fail precisely when intermediate representations blur. ProsQA asks a model to search over a graph of fictional concepts and verify a claim, holding several candidate reasoning paths at once. Arithmetic expressions demand exact intermediate results, where confusing 12 with 13 anywhere poisons everything downstream.

The arithmetic result is the cleanest picture of the mechanism. With four operands both models cope; by six, LLaMA's accuracy falls to 41.3 percent while LIMe keeps 71.6, and the values of nearby numbers remain visibly separated in LIMe's layers. Early layers can hold a partial result and let later layers compose it, instead of overwriting it in the shared stream.

Depth tells the same story at scale. Training at 32, 64, and 128 layers, LIMe wins at every depth, and its loss falls faster as depth grows: a 64-layer LIMe beats a 128-layer LLaMA that spends roughly twice the FLOPs and parameters. The deeper the network, the more a single overwritten stream costs.

What the router chose to remember

The learned routing weights are interpretable on their own. Layers 2 through 4 lean heavily on the embedding buffer, consistent with early attention handling local and morphological structure. Middle layers borrow their immediate predecessors as auxiliary memory, widening the feature subspace heads can manipulate. And the final layers reach all the way back to the first buffers before predicting, revisiting raw token identity one last time.

The ablation ranks the ingredients cleanly. A uniform average of buffers is exactly as bad as no router at all, so learning where to look is the substance, not the access itself. Restricting learned routing to recent layers gives up most of the gain, while restricting it to the earliest layers keeps nearly all of it. The model's favorite memory is its oldest.

The honest costs: a naive implementation grows router terms quadratically with depth, and sharded pipeline-parallel training would need to communicate buffers between stages. Both have workable mitigations, and the restricted-window results above suggest cheap approximations that keep most of the benefit. What we take away is simpler: the memory was already there, sitting in the buffers. We only had to let attention read it.