# Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning

arXiv: https://arxiv.org/abs/2604.21999 · HTML: https://mlsanity.com/paper/2604.21999/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 24 references that are in the map.
- First posted on arXiv 23 Apr 2026.

## Builds on (most cited first, 19 of 19)

- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach](https://mlsanity.com/paper/2502.05171/) (2025)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention](https://mlsanity.com/paper/2404.07143/) (2024)
- [Hierarchical Reasoning Model](https://mlsanity.com/paper/2506.21734/) (2025)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation](https://mlsanity.com/paper/2507.10524/) (2025)
- [The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers](https://mlsanity.com/paper/2108.12284/) (2021)
- [Memory Transformer](https://mlsanity.com/paper/2006.11527/) (2020)
- [Memformer: A Memory-Augmented Transformer for Sequence Modeling](https://mlsanity.com/paper/2010.06891/) (2020)
- [Stronger Normalization-Free Transformers](https://mlsanity.com/paper/2512.10938/) (2025)
- [Universal Reasoning Model](https://mlsanity.com/paper/2512.14693/) (2025)
- [Adaptive Loops and Memory in Transformers: Think Harder or Know More?](https://mlsanity.com/paper/2603.08391/) (2026)
- [MeSH: Memory-as-State-Highways for Recursive Transformers](https://mlsanity.com/paper/2510.07739/) (2025)

## Built on by (most cited first, 5 of 5)

- [Solve the Loop: Attractor Models for Language and Reasoning](https://mlsanity.com/paper/2605.12466/) (2026)
- [Loop the Loopies!](https://mlsanity.com/paper/2607.16051/) (2026)
- [Chain-of-Thought and Compressed Looped Transformers: A Memory-Budget Separation](https://mlsanity.com/paper/2605.30757/) (2026)
- [Repeated Shared Access Enables Grokking, but Edit Propagation Depends on an Addressable Memory](https://mlsanity.com/paper/2606.20737/) (2026)
- [Random Recursive Models](https://mlsanity.com/paper/2610.00541/) (2026)

## Abstract (arXiv)

We study learned memory tokens as a computational scratchpad for a single-block Universal Transformer with Adaptive Computation Time (ACT) on Sudoku-Extreme, a combinatorial reasoning benchmark. Memory tokens are empirically necessary: no configuration without them reaches non-trivial performance. The optimal count has a sharp lower threshold (T=0 always fails, T=8 reliably succeeds) followed by a stable plateau (T=8-32, 57.4% +/- 0.7% exact-match) and a dilution boundary at T=64. Under halt-side pressure (lambda warmup), mean halt drops monotonically with memory size across the plateau (from 11.6 at T=8 to 8.3 at T=64), showing that memory tokens and ponder depth substitute as resources at fixed accuracy. We also identify a router initialization trap that causes the majority of training runs to fail: both default zero-bias and Graves' recommended positive bias settle into a shallow halt equilibrium the model cannot escape. Inverting the bias to -3 ("deep start") eliminates the failure mode, and ablation shows the trap is inherent to ACT initialization rather than an artifact of our architecture. With reliable training, ACT yields an order of magnitude lower seed variance than fixed-depth processing (+/-0.7 vs +/-9.3 pp); lambda warmup recovers 34% of compute at matched accuracy; and attention heads specialize into memory readers, constraint propagators, and integrators across recursive depth. Code: https://github.com/che-shr-cat/utm-jax.
