pith. machine review for the scientific record. sign in

arxiv: 2604.21999 · v3 · submitted 2026-04-23 · 💻 cs.LG · cs.AI· cs.CL

Recognition: unknown

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

Grigory Sapunov

Authors on Pith no claims yet

Pith reviewed 2026-05-09 22:42 UTC · model grok-4.3

classification 💻 cs.LG cs.AIcs.CL
keywords universal transformeradaptive computation timememory tokenscombinatorial reasoningsudokurecursive reasoningdepth-state trade-offattention specialization
0
0 comments X

The pith

Universal Transformers require memory tokens to achieve non-trivial performance on combinatorial reasoning tasks.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper studies how learned memory tokens function as a scratchpad inside a single-block Universal Transformer equipped with Adaptive Computation Time when solving Sudoku-Extreme puzzles. Without any such tokens the model reaches only trivial performance, yet eight or more tokens produce a stable accuracy plateau near 57 percent exact match. The tokens and the model's pondering depth act as interchangeable resources, so larger memory counts reduce the average number of computation steps required while holding accuracy fixed. Proper initialization of the halting router is also shown to be essential for escaping shallow equilibria and achieving low variance across training seeds.

Core claim

Learned memory tokens serve as an explicit computational scratchpad that a single-block Universal Transformer with Adaptive Computation Time needs to solve Sudoku-Extreme. Configurations lacking these tokens fail to reach non-trivial accuracy, while counts from eight to thirty-two tokens deliver consistent 57.4 percent exact-match performance and permit the model to reduce its mean halt steps from 11.6 down to 8.3 as memory size grows. This establishes a direct trade-off between state capacity and recursion depth at constant accuracy, and it requires inverting the router bias to a deep-start value to avoid the shallow-halt trap inherent to Adaptive Computation Time training.

What carries the argument

Learned memory tokens used as a scratchpad that trades off against ponder depth inside Adaptive Computation Time halting.

Load-bearing premise

The observed necessity of memory tokens and their substitution for depth are not artifacts of the Sudoku-Extreme benchmark or the single-block Universal Transformer architecture.

What would settle it

A run of the same architecture on Sudoku-Extreme that reaches high accuracy with zero memory tokens, or one in which increasing memory size fails to reduce mean halt steps at matched accuracy.

Figures

Figures reproduced from arXiv: 2604.21999 by Grigory Sapunov.

Figure 1
Figure 1. Figure 1: UTM architecture. Sequence tokens (orange) and memory tokens (purple) receive separate [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Standard initialization (left) shows extreme seed sensitivity—same architecture, different [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Memory-token curve with deep-start initialization. [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Training trajectories for the four runs of Table 4 ( [PITH_FULL_IMAGE:figures/full_fig_p007_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Attention maps at steps 0, 9, and 17 (head-averaged, [PITH_FULL_IMAGE:figures/full_fig_p008_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Per-head attention at step 17 (T=16). Each head’s s→m fraction is shown. H4 (0.34) and H1 (0.24) are memory-focused; H2 and H6 (0.01) are pure puzzle-constraint heads with periodic S→S structure; H5 (0.21) mixes column-attention stripes with memory reading. See text for per￾head analysis. The model memorizes perfectly but achieves only 6–8% on unseen puzzles, even with TRM￾matched regularization. This sugg… view at source ↗
Figure 7
Figure 7. Figure 7: Extended inference (64 steps, trained on 18). The lambda-warmup model peaks at step 36 [PITH_FULL_IMAGE:figures/full_fig_p009_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Step-by-step puzzle solving (T=16, λ=0.001+warmup). The model progressively fills cells across ponder steps: 31/81 → 53 → 60 → 74 → SOLVED. Green = correct, red = error. 17 [PITH_FULL_IMAGE:figures/full_fig_p017_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: A puzzle the model fails to solve. Despite 18 ponder steps, the model reaches only 43/81 [PITH_FULL_IMAGE:figures/full_fig_p018_9.png] view at source ↗
read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The paper empirically studies learned memory tokens as a scratchpad in a single-block Universal Transformer with Adaptive Computation Time (ACT) on the Sudoku-Extreme combinatorial reasoning benchmark. It claims memory tokens are necessary for non-trivial performance (T=0 always fails), with a sharp threshold at T=8, a stable plateau at T=8-32 yielding 57.4% +/- 0.7% exact-match accuracy, and dilution beyond T=64; under lambda warmup, mean halt steps decrease monotonically with memory size, indicating depth-state substitution. It also identifies a router bias initialization trap (default/positive biases fail, -3 succeeds), reports lower seed variance with ACT, 34% compute recovery via warmup, and attention head specialization.

Significance. If the central empirical results hold under fuller controls, the work usefully demonstrates that explicit memory tokens and ponder depth can trade off as resources in adaptive recursive reasoning, with practical guidance on ACT initialization and evidence of head specialization. The open-sourced code is a clear strength for reproducibility. The findings are specific to one benchmark and architecture but provide falsifiable observations on resource substitution that could inform broader transformer design.

major comments (2)
  1. [Abstract / Experimental results on memory token counts] Abstract and experimental results on memory token counts: the sharp lower threshold (T=8 succeeds, T=0 fails) and stable plateau (T=8-32 at 57.4% +/- 0.7%) are presented as key findings, but the range selection is post-hoc without pre-specified criteria, full reporting of all tested T values, or statistical tests for stability and the dilution boundary at T=64. This is load-bearing for the necessity and trade-off claims.
  2. [Router initialization and halt-side pressure experiments] Router bias initialization and lambda warmup experiments: the inversion to bias=-3 is shown via ablation to eliminate the shallow halt trap, but the manuscript provides no systematic search or sensitivity analysis around this specific value (a free parameter), nor confirms the depth-state substitution holds independently of the warmup schedule. This affects the reliability of the reported trade-off and training success.
minor comments (2)
  1. [Abstract / Results] The variance reduction claim (+/-0.7 vs +/-9.3 pp) and 34% compute recovery are stated numerically but should be explicitly tied to a table or figure for verification.
  2. Clarify the exact definition of 'compute' recovered by lambda warmup and the metric for 'exact-match' accuracy to aid replication.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive and detailed feedback. We address each major comment below, indicating where we agree that additional reporting or experiments are warranted and outlining the specific revisions we will incorporate.

read point-by-point responses

  1. Referee: Abstract and experimental results on memory token counts: the sharp lower threshold (T=8 succeeds, T=0 fails) and stable plateau (T=8-32 at 57.4% +/- 0.7%) are presented as key findings, but the range selection is post-hoc without pre-specified criteria, full reporting of all tested T values, or statistical tests for stability and the dilution boundary at T=64. This is load-bearing for the necessity and trade-off claims.

    Authors: We agree that fuller reporting of the memory token experiments would improve transparency and strengthen the claims. The manuscript highlighted results for T=0, 8, 16, 32, and 64 as representative of the observed threshold, plateau, and dilution effects, but did not include the complete set of tested values or formal statistical comparisons. In the revised manuscript we will add a dedicated table reporting exact-match accuracy (mean and standard deviation across seeds) for all T values we evaluated (including 1, 2, 4, 16, 64, and 128). We will also include a brief description of the exploratory process used to select the reported range and apply simple statistical tests (pairwise t-tests with correction) to quantify the significance of the threshold at T=8 and the stability of the T=8–32 plateau relative to the drop at T=64. These additions directly address the post-hoc concern while preserving the original empirical observations. revision: yes

  2. Referee: Router bias initialization and lambda warmup experiments: the inversion to bias=-3 is shown via ablation to eliminate the shallow halt trap, but the manuscript provides no systematic search or sensitivity analysis around this specific value (a free parameter), nor confirms the depth-state substitution holds independently of the warmup schedule. This affects the reliability of the reported trade-off and training success.

    Authors: The referee correctly identifies that the manuscript demonstrates the effectiveness of bias=-3 via a limited ablation but does not provide a broader sensitivity analysis or isolate the depth-state substitution from the lambda warmup schedule. In the revision we will expand the router-bias section with a sensitivity sweep over initial bias values from -5 to +1 (in integer steps), reporting both training success rate and final accuracy for each setting. We will also add a control experiment that trains with a fixed (non-warmup) lambda schedule and replots mean halt steps versus memory size; this will confirm whether the observed monotonic substitution between memory tokens and ponder depth persists without the warmup schedule. These results will be presented in the revised experimental section. revision: yes

Circularity Check

0 steps flagged

No circularity: purely empirical measurements with no derivations or self-referential fits

full rationale

The paper reports experimental results on memory tokens, ACT halt behavior, and depth-state trade-offs using direct training runs, ablations, and accuracy measurements on Sudoku-Extreme. All key findings (necessity of T>=8, monotonic halt reduction, initialization trap) are obtained from observed data rather than any model equation, fitted parameter, or self-citation chain. No load-bearing step reduces to its own inputs by construction, and external citations (e.g., Graves) are not used to justify uniqueness or ansatzes. The work is self-contained as an empirical study.

Axiom & Free-Parameter Ledger

3 free parameters · 1 axioms · 0 invented entities

The central claims rest on empirical hyperparameter choices (memory token count, lambda warmup schedule, router bias of -3) and the assumption that Sudoku-Extreme is a faithful test of recursive combinatorial reasoning; no new mathematical axioms or invented entities are introduced.

free parameters (3)
  • memory token count T
    Empirically swept; optimal plateau identified post-hoc at T=8-32
  • router bias initialization
    Set to -3 after observing failure of default and Graves positive bias
  • lambda warmup schedule
    Used to apply halt-side pressure; exact schedule not detailed in abstract
axioms (1)
  • domain assumption Sudoku-Extreme requires recursive constraint propagation that benefits from external scratchpad state
    Invoked to interpret why memory tokens are necessary

pith-pipeline@v0.9.0 · 5575 in / 1448 out tokens · 39136 ms · 2026-05-09T22:42:15.624348+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

24 extracted references · 19 canonical work pages · 6 internal anchors

  1. [1]

    Pondernet: Learning to ponder

    A. Banino, J. Balaguer, and C. Blundell. PonderNet: Learning to Ponder. arXiv:2107.05407, 2021. https://arxiv.org/abs/2107.05407

    work page arXiv 2021

  2. [2]

    Advances in Neural Information Processing Systems (NeurIPS) , year =

    A. Bulatov, Y. Kuratov, and M. Burtsev. Recurrent Memory Transformer. In NeurIPS, 2022. https://arxiv.org/abs/2207.06881

    work page arXiv 2022

  3. [3]

    arXiv preprint arXiv:2006.11527 , year=

    M. Burtsev et al. Memory Transformer. arXiv:2006.11527, 2020. https://arxiv.org/abs/2006.11527

    work page arXiv 2006

  4. [4]

    Csord \'a s, K

    R. Csord \'a s, K. Irie, and J. Schmidhuber. The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers. In EMNLP, 2021. https://aclanthology.org/2021.emnlp-main.49/

    2021

  5. [5]

    Darcet et al

    T. Darcet et al. Vision Transformers Need Registers. In ICLR, 2024. https://openreview.net/forum?id=2dnO3LLiJ1

    2024

  6. [6]

    Universal Transformers

    M. Dehghani et al. Universal Transformers. In ICLR, 2019. https://arxiv.org/abs/1807.03819

    work page internal anchor Pith review arXiv 2019

  7. [7]

    A. Graves. Adaptive Computation Time for Recurrent Neural Networks. arXiv:1603.08983, 2016. https://arxiv.org/abs/1603.08983

    work page internal anchor Pith review arXiv 2016

  8. [8]

    URL https: //arxiv.org/abs/2510.04871

    A. Jolicoeur-Martineau. Less is More: Recursive Reasoning with Tiny Networks. arXiv:2510.04871, 2025. https://arxiv.org/abs/2510.04871

    work page arXiv 2025

  9. [9]

    K. Jordan. Muon: An optimizer for hidden layers. https://kellerjordan.github.io/posts/muon/, 2024

    2024

  10. [10]

    Reasoning with latent thoughts: On the power of looped transformers

    N. Saunshi et al. Reasoning with Latent Thoughts: On the Power of Looped Transformers. arXiv:2502.17416, 2025. https://arxiv.org/abs/2502.17416

    work page arXiv 2025

  11. [11]

    Sudoku-Extreme Dataset

    Sapient Intelligence . Sudoku-Extreme Dataset. https://huggingface.co/datasets/sapientinc/sudoku-extreme, 2025

    2025

  12. [12]

    M. Chen, T. Lu, J. Zhu, M. Sun, and Z. Liu. Stronger Normalization-Free Transformers. arXiv:2512.10938, 2025. https://arxiv.org/abs/2512.10938

    work page arXiv 2025

  13. [13]

    Z. Gao, L. Chen, Y. Xiao, H. Xing, R. Tao, H. Luo, J. Zhou, and B. Dai. Universal Reasoning Model. arXiv:2512.14693, 2025. https://arxiv.org/abs/2512.14693

    work page arXiv 2025

  14. [14]

    G. Wang, J. Li, Y. Sun, X. Chen, C. Liu, Y. Wu, M. Lu, S. Song, and Y. A. Yadkori. Hierarchical Reasoning Model. arXiv:2506.21734, 2025. https://arxiv.org/abs/2506.21734

    work page arXiv 2025

  15. [15]

    S. Bae, Y. Kim, R. Bayat, S. Kim, J. Ha, T. Schuster, A. Fisch, H. Harutyunyan, Z. Ji, A. Courville, and S.-Y. Yun. Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation. arXiv:2507.10524, 2025. https://arxiv.org/abs/2507.10524

    work page arXiv 2025

  16. [16]

    Titans: Learning to Memorize at Test Time

    A. Behrouz, P. Zhong, and V. Mirrokni. Titans: Learning to Memorize at Test Time. arXiv:2501.00663, 2025. https://arxiv.org/abs/2501.00663

    work page internal anchor Pith review arXiv 2025

  17. [17]

    M. Frey, B. Shomali, A. H. Bashir, D. Berghaus, J. Koehler, and M. Ali. Adaptive Loops and Memory in Transformers: Think Harder or Know More? arXiv:2603.08391, 2026. https://arxiv.org/abs/2603.08391

    work page arXiv 2026

  18. [18]

    Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

    J. Geiping, S. McLeish, N. Jain, J. Kirchenbauer, S. Singh, B. R. Bartoldson, B. Kailkhura, A. Bhatele, and T. Goldstein. Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. arXiv:2502.05171, 2025. https://arxiv.org/abs/2502.05171

    work page internal anchor Pith review arXiv 2025

  19. [19]

    Neural Turing Machines

    A. Graves, G. Wayne, and I. Danihelka. Neural Turing Machines. arXiv:1410.5401, 2014. https://arxiv.org/abs/1410.5401

    work page internal anchor Pith review arXiv 2014

  20. [20]

    Graves, G

    A. Graves, G. Wayne, M. Reynolds, T. Harley, I. Danihelka, A. Grabska-Barwi \'n ska, S. G. Colmenarejo, E. Grefenstette, T. Ramalho, J. Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626):471--476, 2016. https://www.nature.com/articles/nature20101

    2016

  21. [21]

    Leave no context behind: Efficient infinite context transformers with infini-attention

    T. Munkhdalai, M. Faruqui, and S. Gopal. Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention. arXiv:2404.07143, 2024. https://arxiv.org/abs/2404.07143

    work page arXiv 2024

  22. [22]

    Q. Wu, Z. Lan, K. Qian, J. Gu, A. Geramifard, and Z. Yu. Memformer: A Memory-Augmented Transformer for Sequence Modeling. arXiv:2010.06891, 2020. https://arxiv.org/abs/2010.06891

    work page arXiv 2010

  23. [23]

    C. Yu, X. Shu, Y. Wang, Y. Zhang, H. Wu, J. Li, R. Long, Z. Chen, Y. Xu, W. Su, and B. Zheng. MeSH: Memory-as-State-Highways for Recursive Transformers. arXiv:2510.07739, 2026. https://arxiv.org/abs/2510.07739

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  24. [24]

    R.-J. Zhu, T. Peng, T. Cheng, X. Qu, J. Huang, D. Zhu, H. Wang, K. Xue, X. Zhang, Y. Shan, et al. Scaling Latent Reasoning via Looped Language Models (Ouro). arXiv:2510.25741, 2025. https://arxiv.org/abs/2510.25741

    work page arXiv 2025