Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning
Not on the map. Why?
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.
Key facts. Cited by 5 papers in the map (data through 7 Oct 2026). Builds on 19 of its 24 references that are in the map. First posted on arXiv 23 Apr 2026.
Builds on
19 of 24 references.
References · 19
- 2023Vision Transformers Need Registerscited by 844
- 2014Neural Turing Machinescited by 630
- 2018Universal Transformerscited by 613
- 2016Adaptive Computation Time for Recurrent Neural Networkscited by 479
- 2025Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approachcited by 372
- 2024Titans: Learning to Memorize at Test Timecited by 303
- 2025Reasoning with Latent Thoughts: On the Power of Looped Transformerscited by 223
- 2022Recurrent Memory Transformercited by 197
- 2024Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attentioncited by 194
- 2025Hierarchical Reasoning Modelcited by 167
- 2021PonderNet: Learning to Pondercited by 156
- 2025Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computationcited by 120
7 more
- 2021The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformerscited by 117
- 2020Memory Transformercited by 74
- 2020Memformer: A Memory-Augmented Transformer for Sequence Modelingcited by 63
- 2025Stronger Normalization-Free Transformerscited by 13
- 2025Universal Reasoning Modelcited by 11
- 2026Adaptive Loops and Memory in Transformers: Think Harder or Know More?cited by 11
- 2025MeSH: Memory-as-State-Highways for Recursive Transformerscited by 8
Built on by
Citing papers per year: 2026: 5.
5 papers.
- 2026Solve the Loop: Attractor Models for Language and Reasoningcited by 11
- 2026Loop the Loopies!cited by 4
- 2026Chain-of-Thought and Compressed Looped Transformers: A Memory-Budget Separationcited by 2
- 2026Repeated Shared Access Enables Grokking, but Edit Propagation Depends on an Addressable Memorycited by 0
- 2026Random Recursive Modelscited by 0
- 2026Random Recursive Models30 Sep 2026
- 2026Loop the Loopies!17 Jul 2026
- 2026Repeated Shared Access Enables Grokking, but Edit Propagation Depends on an Addressable Memory17 Jun 2026
- 2026Chain-of-Thought and Compressed Looped Transformers: A Memory-Budget Separation29 May 2026
- 2026Solve the Loop: Attractor Models for Language and Reasoning12 May 2026