The Surprising Effectiveness of Shared Memory in Looped Transformers

Authors: Giovanni Monea, Keshav Ramji, Yousef El-Kurdi, Luis A. Lastras, Yoav Artzi, Nathan Godey, Ramón Fernandez Astudillo

---

Summary (Overview)

• Core Contribution: The paper introduces Memory Sharing as an architectural inductive bias for looped Transformers, where later recursions read the KV cache written by the first recursion while keeping only a short local window of their own.
• Two New Models: The Looped Prediction Transformer (LPT) applies memory sharing to full-attention looped models, and the Hybrid Looped Prediction Transformer (H-LPT) applies it to hybrid linear-attention/full-attention models.
• Key Finding: Contrary to expectations, memory sharing improves quality rather than degrading it. LPT and H-LPT outperform standard looped Transformers (LT, LT²) in perplexity and downstream accuracy while using 76–79% less context memory than a standard Transformer.
• Mechanistic Insights: Shared memory acts as a "gradient highway" to the first recursion, and shared/local memory develop distinct representations, with later recursions attending predominantly to shared memory.
• Efficiency Gains: The memory savings translate to 2.0–4.8× more sequences fitting in a GPU batch compared to standard looped models, with nearly identical inference FLOPs.

---

Introduction and Theoretical Foundation

Background

Looped Transformers (LT) repeatedly apply the same layers several times per token, trading computation for quality without adding parameters (Dehghani et al., 2019; Zhu et al., 2025). However, in standard formulations, each recursion writes its own key-value (KV) cache, causing context memory to grow linearly with recursion depth even though parameter count stays fixed.

Motivation

Previous work showed that sharing caches across recursions during decoding can save memory with little quality loss (Geiping et al., 2025; Zhu et al., 2025). This raises a key question: if sharing can work in models never trained for it, what can it do when models are trained for it? The authors investigate memory sharing as an architectural inductive bias from the start of training.

Theoretical Foundation

The paper builds on the Ouro full-stack formulation of looped Transformers (Zhu et al., 2025). Let $M_L = T^{(L)}_{\theta} \circ \cdots \circ T^{(1)}_{\theta}$ denote the shared stack of $L$ Transformer blocks. For a sequence of length $T$ and hidden dimension $d$, the stack is applied $R$ times:

$$h^{(1)}_{in} = \text{emb}(x), \quad h^{(t)} = M_L \left( h^{(t)}_{in} \right), \quad h^{(t+1)}_{in} = N_{rec} \left( h^{(t)} \right), \quad \text{logits}^{(t)} = \mathrm{lm\_head} \left( N_{out}(h^{(t)}) \right), \quad t = 1, \ldots, R$$

where $N_{rec}$ and $N_{out}$ are separately parameterized RMSNorms. The effective depth is $L \cdot R$ with the parameters of a single $L$-layer stack.

For vanilla LT, a query at position $i$ in recursion $t$ reads the causal prefix of its recursion's cache:

$$A^{(t)}_{LT}(i) = C^{(t)}_{\leq i} = \{(k^{(t)}_j, v^{(t)}_j) : j \leq i\}$$

---

Methodology

Memory Sharing Architecture

The first recursion (the state recursion) runs as in LT and writes a full-context cache $C^{(1)}$ at each layer. Every later recursion (a prediction recursion) attends to this shared cache together with its own KV entries using a single softmax. Each prediction recursion's own cache is restricted to a short sliding window of $\omega$ preceding positions:

$$A^{(t)}_{LPT}(i) = \begin{cases} C^{(1)}_{\leq i}, & t = 1 \\ C^{(1)}_{\leq i} \cup C^{(t)}_{[i-\omega, i]}, & t = 2, \ldots, R \end{cases}$$

Model Variants

• LPT: Full-attention looped model with memory sharing
• H-LPT: Hybrid model (Gated DeltaNet + full attention layers) with memory sharing applied to full-attention layers

KV Cache Memory

The inference-time KV cache sizes across $L$ full-attention layers are:

$$M^{LT}_{KV} = L \cdot R \cdot T, \qquad M^{LPT}_{KV} = L [T + (R-1)\omega]$$

At $T = 4{,}096$ with $\omega = 64$ and five recursions, windows increase attended KV entries by only $(R-1)\omega/(RT) = 1.25\%$ while reducing full-attention cache by $4.7\times$.

Experimental Setup

• Model sizes: S (151.9M), M (454.2M), L (1.01B) parameters
• Context length: 4,096 tokens
• Data: FineWeb-Edu, 20B tokens (S, M) and 50B tokens (L)
• Baselines: Standard Transformer (STD), LT, LT², PLT (Parallel Loop Transformer), LOOPIE, DEEPSTD
• Training objectives: Final loss (NLL at final recursion only) and Adaptive loss (Ouro's weighted per-recursion NLL with KL regularization)
• Evaluation: Perplexity on held-out FineWeb-Edu + zero-shot accuracy on ARC-Easy, HellaSwag, PIQA, SciQ, LAMBADA

---

Empirical Validation / Results

Main Results (Table 1, selected highlights)

At the L size (1B parameters), five recursions:

Key results:

• LPT vs. LT: LPT lowers perplexity by 0.38–1.09 nats and raises accuracy by 0.4–4.8 points with 2.0–4.7× less KV memory and under 1% more FLOPs
• H-LPT vs. LT²: H-LPT lowers perplexity by 0.09–0.15 and raises accuracy by 0.8–1.2 points with 4.1–4.2× less context memory
• H-LPT vs. STD: At five recursions, H-LPT lowers perplexity by 1.12–1.82 with 76–79% less context memory

Why Memory Sharing Helps

The authors conduct extensive analysis to understand the quality improvement:

1. Stability with final loss: LPT trains stably with final loss, while LT degrades or collapses at depth without adaptive loss.

2. Window size robustness: LPT outperforms LT at every window size ($\omega = 0$ to $\omega = 4{,}096$), even when prediction recursions keep no local history.

3. Attention allocation: Later recursions attend mostly to shared memory (57–80% of attention mass), with preference increasing in later recursions.

4. Representation divergence: Shared and local keys form distinct clusters in PCA analysis, suggesting they serve different purposes (memory construction vs. prediction refinement).

5. Gradient highway: Shared memory provides a direct gradient path from the final loss to the first recursion. The gradient norm reaching the first recursion falls sharply with depth in LT but is largely recovered in LPT:

$$p_a = \frac{\langle g_a, g^{(1)} \rangle}{\|g^{(1)}\|^2}, \quad a \in \{KV, \text{hidden}\}, \quad p_{KV} + p_{\text{hidden}} = \|g^{(1)}\|^2$$

6. Detach ablation: Detaching shared KV gradients during training hurts learning—perplexity rises at every depth, and deeper models often fail to train entirely.

Statistical Significance

At S size with four seeds, paired t-tests confirm significance:

• LPT vs. LT: $p < 10^{-5}$ at all recursion counts
• H-LPT vs. LT²: $p < 0.002$ at all recursion counts
• LPT vs. PLT: $p < 0.05$ at all recursion counts

GPU Efficiency (Table 8, L size, 5 recursions)

H-LPT fits 3.6–5.0× as many sequences as STD on one GPU and generates 2.1–2.2× tokens/sec at two recursions, 1.4–1.5× at three, and 0.8–0.9× at five.

---

Theoretical and Practical Implications

Theoretical Implications

1. Memory–prediction separation: The results support a division of labor across recursions: the first recursion specializes in building memory while later recursions use it for prediction. This parallels the State-Prediction Separation hypothesis (Monea et al., 2026).

2. Representation similarity not required: Prior work attributed the small cost of cache reuse to similarity of keys/values across recursions. LPT shows that distinct representations can coexist effectively when sharing is trained for.

3. Gradient flow in recurrent depth: The gradient highway mechanism provides a principled explanation for why memory sharing stabilizes deep recursion, addressing a known instability issue in looped models.

Practical Implications

1. Quality–memory frontier: LPT and H-LPT significantly push the quality–memory trade-off frontier for looped models, making them viable for long-context applications.

2. Compute efficiency: H-LPT achieves better quality than a standard Transformer at similar or higher throughput (2–3 recursions) or 81–88% throughput (5 recursions), with substantially less memory.

3. Scalability: The memory savings allow 2.0–4.8× larger batches on the same GPU hardware, potentially translating to better hardware utilization and lower serving costs.

---

Conclusion

The paper demonstrates that shared memory is a useful inductive bias for looped models, not merely a memory-saving technique. The key findings are:

1. Memory Sharing improves quality over full-cache looped baselines across all tested sizes (150M–1B) and recursion counts (2–9).
2. The mechanism involves: (a) later recursions attending predominantly to shared memory, (b) shared and local memory developing specialized representations, and (c) shared memory creating a gradient highway that stabilizes training.
3. Hybrid variants (H-LPT) achieve the best quality–memory–throughput trade-off, outperforming same-size standard Transformers with 76–79% less context memory.

Limitations and Future Work

• Compute-optimality: LPT trades time for quality; at five recursions it costs 4.3–4.8× the decoding FLOPs of STD.
• Fixed recursion count: The model spends the same number of recursions on every token; adaptive per-token recursion selection is future work.
• Scale: Training is limited to 1B parameters on one corpus with 4,096-token contexts.
• Future directions: Increasing the number of state recursions, alternating state and prediction recursions, and more sophisticated memory sharing approaches.

---

Key Takeaways

Takeaway 1: Through Memory Sharing, LPT and especially H-LPT improve the quality–memory frontier for looped models.

Takeaway 2: LPT stays stable and trains best with final loss.

Takeaway 3: Shared memory improves quality across prediction-window sizes, even when later recursions keep no local history.

Takeaway 4: Later recursions attend more to shared memory than to their own caches.

Takeaway 5: Shared and local memory diverge in their local representations, so sharing is not reliant on recursions writing similar memories.

Takeaway 6: The memory sharing mechanism creates a gradient highway from the final loss to the first recursion that facilitates learning.