# The Surprising Effectiveness of Shared Memory in Looped Transformers

> Memory sharing in looped Transformers improves quality and stabilizes training by creating a gradient highway, while cutting context memory by up to 79%.

- **Source:** [arXiv](https://arxiv.org/abs/2610.02383)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/DnKkYH
- **Whiteboard:** https://picx.dev/p/DnKkYH/image

## Summary

# 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)} = \text{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:

| Model | Perplexity ↓ | Accuracy ↑ | KV Memory (×STD) ↓ | FLOPs (×STD) ↓ |
|-------|-------------|------------|---------------------|-----------------|
| STD | 12.27 | 61.4 | 2.25 | 1.32 |
| LT | 11.39 | 63.4 | 5.00 | 4.81 |
| LPT | **10.32** | **68.2** | 1.06 | 4.83 |
| LT² | 10.32 | 67.3 | 1.01 | 3.75 |
| H-LPT | **10.23** | **68.1** | **0.24** | 3.75 |

**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)

| Model | Batch Size | Tokens/sec | Latency (ms) |
|-------|-----------|------------|-------------|
| STD | 108 | 3,162 | 6.7 |
| LT | 21 | 303 | 37.9 |
| LPT | 101 | 578 | 42.3 |
| LT² | 103 | 1,776 | 21.7 |
| H-LPT | **390** | **2,553** | 22.6 |

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.

---

_Markdown view of https://picx.dev/p/DnKkYH, served by PicX — AI-generated visual whiteboard summaries of research papers._