Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning — part 3/3

Cleaned-text code points [31182, 44830). Bibliography omitted here; source.md includes it.

For completeness, we note two adjacent lines of work that share vocabulary with ours but address different problems. Memformer ( Wu et al., 2020 ) , Infini-attention ( Munkhdalai et al., 2024 ) , and Titans ( Behrouz et al., 2025 ) address long-context efficiency through external memory modules updated across sequence segments. The Neural Turing Machine ( Graves et al., 2014 ) and Differentiable Neural Computer ( Graves et al., 2016 ) introduced explicit differentiable memory with content and location-based addressing. These target sequence-length recurrence or explicit data-structure operations rather than depth recurrence within a single forward pass, but inform the design-space discussion below.

### 6.6 Mapping the Design Space

Architectures across the threads above can be characterized along three axes: (i) how the recurring unit is structured, (ii) how persistent state is carried across recursion, and (iii) how computation is allocated. Table 9 maps the closest neighbors. Our work occupies the simplest configuration along each axis: a single transformer layer as recurring unit, sample-independent recursion-updated memory tokens as persistent state, ACT-mediated variable depth as computation allocation, and a single forward pass without deep supervision.

The design-space view clarifies several relationships that surface vocabulary obscures.

TRM and UTM use distinct memory mechanisms despite both involving prepended learnable tokens. TRM’s puzzle-ID embedding ( puzzle_emb in the public code) is concatenated to the input sequence as additional token positions, in the Burtsev et al. (2020) lineage, but functions as a per-puzzle-ID static identifier rather than a recursion-updated scratchpad: it is set once per puzzle and re-injected unchanged at each recursion step. Our memory tokens are sample-independent (shared across all inputs of a given task) and update via attention each ponder step, functioning as a general-purpose recursion-updated scratchpad. TRM additionally maintains its full-sequence-shaped latents y y and z z , summed (not concatenated) into the input embedding, as the actual recursive state—a mechanism we do not use.

MeSH’s mechanism occupies a different point in the design space than ours, despite reaching a similar conclusion. MeSH’s slots are full-sequence-shaped tensors accessed by per-iteration routers that are not weight-tied across iterations. Our memory tokens are distinct prepended positions accessed by standard self-attention with a fully weight-shared block. The two mechanisms address overlapping problems via different architectural choices: MeSH addresses functional specialization across iterations by making the routers iteration-specific; we maintain full weight-tying and address the same underlying issue (insufficient bandwidth in the residual stream) by adding parallel addressable storage. Whether iteration-specific routing and addressable storage are complementary or substitutable is an open empirical question.

URM and UTM are closer architecturally to each other than to TRM or HRM. Both use the explicit UT+ACT framing; both rely on the residual stream rather than per-position latent state for working memory. URM substitutes ConvSwiGLU and truncated backpropagation for our memory tokens and single-pass training. The two architectures are evaluated under different training regimes—URM follows the HRM/TRM 1K-example small-data protocol with heavy augmentation, while our 57% is reported under full training—so a direct head-to-head comparison would be misleading; isolating the contribution of each architectural axis under a matched regime is left to future work.

Ouro shows that loop depth alone scales without explicit memory augmentation, but identifies the same manipulation-vs-capacity tradeoff. Ouro reports that recurrence does not increase per-parameter knowledge storage but enhances knowledge manipulation. This parallels Frey et al.’s explicit dissociation of looping (manipulation) from memory (capacity), and is consistent with our finding that memory tokens become necessary precisely when the architecture is most minimal—exactly the regime where the residual stream’s capacity is most constrained.

Our “memory tokens are necessary” finding is specific to the minimal configuration. TRM solves Sudoku-Extreme without our memory-token mechanism but with full-sequence per-position summed latents and deep supervision. URM solves Sudoku-Extreme without explicit persistent state but with ConvSwiGLU and TBPTL. MeSH improves recursive language modeling without standard memory tokens but with externalized iteration-specific routing. Frey et al. recover commonsense performance via key-value memory banks. Ouro relies on loop depth alone. The finding is therefore not that memory tokens are universally necessary for recursive computation, but that something must compensate for the residual stream’s narrow bandwidth in a depth-recurrent architecture—and across the works above, every architecture that succeeds adds some such mechanism. Memory tokens are one such compensation; per-position summed latents are another; iteration-specific external buffers are a third; gated key-value retrieval is a fourth; convolutional sequence mixing is a fifth; sufficient parameter count combined with massive pretraining is a sixth. Our work establishes the minimal mechanism along this axis on the smallest architecture and cleanest task; characterizing the full equivalence class is open.

### 6.7 What the Convergence Suggests

The threads above started from different motivations—adaptive computation budgets, long-context memory, biologically-inspired recursion, parameter-efficient pretraining at scale—and have arrived in 2025–2026 at architectures that share a substantial core: weight-shared recurrence, persistent state across recursion, and (often) adaptive halting. This convergence is now visible across scales, from minimal architectures on combinatorial tasks (this work) to language models pretrained on trillions of tokens ( Zhu et al., 2025 ; Geiping et al., 2025 ; Yu et al., 2026 ) .

A useful framing emerges from this convergence: looping enhances the model’s ability to manipulate information through iterative refinement, while memory enhances its capacity to store information that the manipulation operates on. Frey et al. (2026) articulate this dissociation explicitly; Ouro corroborates it independently with the empirical observation that recurrence does not increase per-parameter knowledge storage but dramatically improves multi-hop reasoning. Under this framing, our finding that memory tokens are necessary in the minimal UT setting—where residual-stream capacity is most constrained—is exactly what one would predict: the minimal residual stream cannot manipulate knowledge it has no capacity to store. The bandwidth of working memory and the depth of computation are complementary architectural resources, and the right ratio depends on task structure.

This convergence motivates several testable predictions:

1. Bandwidth-allocation scaling. Wider models (more residual capacity) should require fewer memory tokens, while deeper recursion (more state to maintain across steps) should require more—and varying these axes systematically should trace iso-performance contours rather than independent optima.

2. Additive vs. substitutive gains. Since TRM solves Sudoku-Extreme via full-sequence summed latents and URM via convolutional sequence mixing, augmenting either with explicit recursion-updated memory tokens should yield additive rather than substitutive gains; if the gains are instead substitutive, the bandwidth framing must be refined.

3. Dissociability of halting and memory. Within the recursive-transformer family halting and memory are dissociable mechanisms—Ouro and PonderNet provide adaptive halting without explicit memory augmentation, while MeSH provides explicit memory without adaptive halting (fixed N loop N_{\mathrm{loop}} )—so we expect that combining them deliberately (rather than treating them as solutions to the same problem) is the architecturally correct move.

Characterizing the full equivalence class of bandwidth-augmenting mechanisms across tasks and scales, and verifying or falsifying these predictions, is the natural next direction for the field.

## 7 Future Work

Our findings open several directions we consider most promising: (1) Scaling and width revisited : our 3.2M model is half the size of TRM (7M). A hidden=768 model ( ∼ 7 {\sim}7 M params) would enable direct comparison. Notably, all prior width experiments (hidden=768) used the default initialization and uniformly failed (2–5% EM)—it is unknown whether deep-start initialization would rescue these, which would reframe the “width-variance trap” observed in early experiments as another manifestation of the router initialization issue. (2) Muon optimizer : Newton-Schulz-based optimizers ( Jordan, 2024 ) have shown promise for weight-shared architectures; combining Muon with deep-start initialization and lambda warmup is unexplored. (3) TRM-style training protocol : our small-dataset results (6–8% generalization) use a single-pass encoder; adapting TRM’s iterative answer-refinement mechanism to the UT framework could improve algorithmic generalization. (4) Multi-task evaluation : testing on maze navigation, ARC-AGI, or formal logic to determine whether the T = 8 T{=}8 threshold and attention specialization patterns are task-universal or Sudoku-specific. (5) Deeper attention analysis : per-head, per-difficulty probing to understand whether memory readers and constraint propagators emerge consistently across tasks and scales.

## 8 Limitations

(1) Single task : all results on Sudoku-Extreme; thresholds are likely task-specific. (2) Architecture-specific : HRM, TRM, and URM all solve Sudoku without memory tokens, via different mechanisms (hierarchical z L , z H z_{L},z_{H} state; per-position y , z y,z latents with deep supervision; ConvSwiGLU + TBPTL, respectively); our findings apply to the single-block UT. A parameter-matched comparison (7M) is planned. (3) Algorithmic generalization : the single-block UT memorizes but does not generalize (6–8% vs TRM’s 87.4%) on the 1K-example protocol, likely reflecting architectural differences in learning mechanisms rather than capacity alone (3.2M vs TRM’s 7M params). (4) T = 64 T{=}64 seed sensitivity : deep-start eliminates sensitivity in the plateau but not at the dilution boundary. (5) Fixed-depth variance : fixed-depth processing shows high seed variance ( ± 9.3 % \pm 9.3\% across 3 seeds) compared to ACT ( ± 0.7 % \pm 0.7\% ), likely due to reliance on a single step’s representation. (6) Evaluation granularity : eval EM is computed on 12.8K-sample batches (SE ± \pm 0.4pp), well below the inter-seed variance ( ± \pm 0.7pp); full 423K test-set evaluation would not materially change the reported means.

## 9 Conclusion

We demonstrate that learned memory tokens are empirically necessary for a single-block Universal Transformer to solve combinatorial reasoning tasks. Without them, no configuration succeeds—regardless of initialization, seed, ponder steps, or use of ACT. With 8 or more tokens ( ∼ 1 {\sim}1 per 10 puzzle cells), the model reliably achieves 57% exact-match on Sudoku-Extreme, with a stable plateau through T = 32 T{=}32 .

A key methodological finding enables these results: in our experiments, standard ACT initialization creates a degenerate p ≈ 0.5 p\approx 0.5 equilibrium causing > > 70% of runs to fail. This trap—confirmed inherent to ACT via normalization ablation—is fixable with a single line of code.

ACT provides more consistent results than fixed-depth processing ( 56.9 ± 0.7 % 56.9\pm 0.7\% vs 53.4 ± 9.3 % 53.4\pm 9.3\% ), and lambda warmup achieves matching accuracy ( 57.0 ± 1.1 % 57.0\pm 1.1\% ) using 34% fewer ponder steps—genuine compute efficiency without quality loss. Across the T = 8 T{=}8 – 32 32 plateau, memory tokens and ponder depth substitute at fixed accuracy. Under halt-side pressure ( λ = 0.001 \lambda{=}0.001 +warmup, S = 0 S{=}0 ), the halt-vs- T T curve is monotonically decreasing ( 11.60 → 11.50 → 10.29 → 8.25 11.60\to 11.50\to 10.29\to 8.25 for T ∈ { 8 , 16 , 32 , 64 } T\in\{8,16,32,64\} ) at near-constant accuracy on the plateau, with a 2pp drop at the dilution boundary ( T = 64 T{=}64 ). This is the canonical depth-state trade-off: memory tokens and recursive depth function as substitutable resources for combinatorial reasoning, with the dilution boundary intrinsic to the architecture rather than an artifact of saturated halting.

## Appendix A Puzzle Solving Visualization

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .