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

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

The model memorizes perfectly but achieves only 6–8% on unseen puzzles, even with TRM-matched regularization. This suggests that TRM’s autoregressive answer-improvement mechanism may be better suited for algorithmic generalization, though the architectures differ in multiple ways and isolating the causal factor remains future work.

## 5 Making ACT Efficient

### 5.1 Lambda Warmup

Direct application of the ponder penalty collapses halting even with deep start. Lambda warmup resolves this: the model establishes deep processing during the warmup phase, then the penalty compresses computation to halt ≈ 11 \approx 11 . Across 3 seeds, warmup achieves 57.0 ± 1.1 % 57.0\pm 1.1\% EM—matching λ = 0 \lambda{=}0 quality ( 56.9 ± 0.7 % 56.9\pm 0.7\% ) with 34% fewer steps.

### 5.2 Inference Beyond Trained Depth

Models generalize to more ponder steps than trained (step embeddings wrap modularly). Running the lambda-warmup model for 36 steps (2 × \times training depth) yields 66% EM, up from 52% at step 17—a 14pp gain from inference-time compute alone, with no retraining. Beyond 36 steps, quality plateaus and slowly degrades (64% at step 48, 64% at step 63), but never crashes—the model degrades gracefully. The sweet spot is ∼ 2 × {\sim}2\times training depth.

Other models improve modestly (+4–6pp) and plateau by step 28. The practical implication: one can train with aggressive lambda for efficiency (halting at ∼ 11 {\sim}11 steps), then run inference at 2 × 2\times depth when accuracy matters, recovering quality that the penalty compressed away.

## 6 Related Work and Architectural Convergence

The architecture we study sits at the intersection of three threads of research that developed largely independently and have re-engaged in 2025–2026 around a similar set of commitments: applying a single weight-shared block recursively, with adaptive halting, and some form of persistent state across recursion. We trace each thread, then map the design space they collectively occupy.

### 6.1 Universal Transformers and Adaptive Computation Time

Adaptive Computation Time was introduced by Graves (2016) for recurrent neural networks as a mechanism for the network to learn how much computation to spend on each input. Each step produces a halting probability via a sigmoid output; cumulative halting probabilities trigger early termination, and a remainder term provides the gradient path that lets the network learn when to stop.

Dehghani et al. (2019) combined ACT with the Transformer architecture in the Universal Transformer (UT), applying a single weight-shared block iteratively up to a maximum depth. UTs decouple parameter count from depth and provide a path to architectures with theoretical universality properties. Empirically, however, ACT in UTs has shown mixed results—for example, Csordás et al. (2021) implement Universal Transformers as “simply Transformers with shared weights between layers, without adaptive computation time” (§2.2), reflecting a broader trend of dropping ACT from UT variants in practice. Our findings on the router initialization trap (§ 3 ) suggest that some of these difficulties may reflect optimization pathologies in standard initialization rather than fundamental limitations of the architecture.

PonderNet ( Banino et al., 2021 ) reformulates ACT’s halting policy as a probabilistic model with a geometric distribution over halting steps, motivated by the observation that “ACT is notably unstable and sensitive to the choice of a hyper-parameter that trades-off accuracy and computation cost. Additionally, the gradient for the cost of computation can only back-propagate through the last computational step, leading to a biased estimation of the gradient.” PonderNet’s reformulation provides unbiased gradient estimates while remaining fully differentiable.

Saunshi et al. (2025) provide theoretical and empirical evidence that “for many synthetic reasoning problems like addition, p p -hop induction, and math problems, a k k -layer transformer looped L L times nearly matches the performance of a k ​ L kL -layer non-looped model, and is significantly better than a k k -layer model.” This supports the case for weight-shared depth recurrence as a general principle, of which our minimal architecture is one instantiation.

### 6.2 Memory Tokens and Registers

Memory tokens were introduced by Burtsev et al. (2020) in the Memory Transformer, where a small set of learnable vectors are prepended to the input sequence to “store non-local representations” and create a “memory bottleneck for the global information”—a learned global scratchpad attended to by all sequence positions through standard self-attention. Bulatov et al. (2022) extended the mechanism for long-context recurrence in the Recurrent Memory Transformer (RMT), where memory tokens are passed across sequence segments to enable processing beyond the model’s nominal context window.

Darcet et al. (2024) independently rediscovered the same primitive for vision transformers under the name registers , motivated by the observation that “high-norm tokens appearing during inference primarily in low-informative background areas of images” are repurposed by the model for internal computations. Their fix was to append additional tokens to the token sequence—independent of the input image—that the model can learn to use as registers, which absorbs the high-norm pathology and restores clean attention maps. The two genealogies—memory tokens for long-context language modeling, registers for vision transformer interpretability—describe the same architectural mechanism: prepended learnable tokens that participate fully in attention but are not tied to any input position.

Our work applies this mechanism inside the depth-recurrent loop of a UT. Memory tokens in our architecture persist across ponder steps within a single forward pass and serve as recursion-updated working memory rather than long-context summary or interpretability scaffolding.

### 6.3 Recursive Reasoning on Combinatorial Benchmarks

A recent line of work has focused on weight-shared recursive computation under the name recursive reasoning , motivated by performance on combinatorial puzzle benchmarks (Sudoku, mazes, ARC-AGI).

HRM ( Wang et al., 2025 ) introduced two interdependent recurrent modules—a high-level module for “slow, abstract planning” and a low-level module for “rapid, detailed computations”—both implemented as encoder-only Transformer blocks with identical architectures (4 layers each, per Jolicoeur-Martineau, 2025 ), with adaptive halting via Q-learning and deep supervision across up to 16 segments. HRM achieves strong small-data generalization with ∼ 1 {\sim}1 K training examples plus heavy augmentation. Its Related Work explicitly positions the architecture within the Universal Transformer / ACT lineage, noting that “Like earlier neural reasoning algorithms including the Universal Transformer, HRM is computationally universal when given sufficient memory and time constraints” (§6).

TRM ( Jolicoeur-Martineau, 2025 ) simplifies HRM’s two-network design to “a single network with only 2 layers.” The architecture maintains two recursive latent variables: a proposed solution y y and a latent reasoning feature z z (the paper explicitly renames HRM’s z H z_{H} and z L z_{L} for the TRM algorithm: “there is simply an input x x , a proposed solution y y (previously called z H z_{H} ), and a latent reasoning feature z z (previously called z L z_{L} )”). The MLP-Mixer variant (TRM-MLP, 5M parameters) reaches 87.4% on Sudoku-Extreme; the attention variant (TRM-Att, 7M parameters) reaches 74.7% on Sudoku-Extreme and 85.3% on Maze-Hard. TRM trains via deep supervision combined with full back-propagation through the recursion (rather than HRM’s 1-step gradient approximation) and uses a simplified halting mechanism that drops the continue Q-value entirely.

URM ( Gao et al., 2025 ) explicitly identifies its architecture as a Universal Transformer variant. Its §2.2 (“Universal Transformer”) opens by attributing the architecture to Dehghani et al. (2019) : “The Universal Transformer (UT) extends the standard Transformer by introducing recurrent computation over depth. Instead of stacking L L distinct layers, the UT applies a single transition block repeatedly to refine token representations.’’ URM applies UT+ACT in an encoder-style (non-causal) design with a ConvSwiGLU MLP and Truncated Backpropagation Through Loops (TBPTL), reaching 77.6% on Sudoku and 53.8% pass@1 on ARC-AGI-1. 1 1 1 The URM paper describes the architecture as “decoder-only,” but the public implementation uses bidirectional self-attention ( causal=False in models/urm/urm.py ); we describe it accordingly. This matches the HRM and TRM lineage, which likewise refine a fixed-length puzzle sequence rather than generate autoregressively. URM positions its contribution as identifying which mechanisms within the UT family drive performance, finding that “recurrent inductive bias” and the ConvSwiGLU nonlinearity matter more than elaborate architectural designs.

Across these three works, distinct architectural commitments are added to the underlying UT-with-ACT skeleton: HRM contributes hierarchical state, deep supervision, and the 1-step gradient approximation; TRM contributes the single-network simplification, full-recursion backpropagation, and the renaming-and-clarification of HRM’s two latents into a clean ( x x , y y , z z ) structure; URM contributes ConvSwiGLU, TBPTL, and explicit positioning within the UT lineage. Our work occupies the simpler corner of this space—single block, no per-position summed latent state, single-pass training—and isolates what this minimal configuration requires to function.

### 6.4 Recursive Transformers at Language-Modeling Scale

A complementary line of recent work has scaled weight-shared recursive computation to language-modeling pretraining, providing convergent evidence about the role of memory mechanisms.

Geiping et al. (2025) (Huginn) study “a novel language model architecture that is capable of scaling test-time computation by implicitly reasoning in latent space.” Their model “iterat[es] a recurrent block, thereby unrolling to arbitrary depth at test-time,” and is scaled to 3.5 billion parameters and 800 billion tokens. They show the model can improve reasoning performance with additional inference-time compute “up to a computation load equivalent to 50 billion parameters.” This is the canonical LM-scale recurrent-depth pretraining work and demonstrates that the depth-as-third-axis premise of UT/ACT survives into the modern LM regime.

Zhu et al. (2025) (Ouro) extend this thread with “a family of pre-trained Looped Language Models (LoopLM) that …build reasoning into the pre-training phase through (i) iterative computation in latent space, (ii) an entropy-regularized objective for learned depth allocation, and (iii) scaling to 7.7T tokens.” Ouro’s halting mechanism is explicitly linked to PonderNet: the authors describe an exit gate that runs in parallel with the LM head at each step and aligns with adaptive-computation methods like PonderNet, which also optimize an ELBO objective for dynamic halting. At 1.4B and 2.6B parameters, Ouro models match the results of up to 12B SOTA LLMs across a wide range of benchmarks. Importantly, Ouro reports that “recurrence does not increase raw knowledge storage (approximately 2 bits per parameter for both looped and non-looped models) but dramatically enhances knowledge manipulation capabilities on tasks requiring fact composition and multi-hop reasoning.” Ouro carries information through the hidden state alone, without any auxiliary memory mechanism.

Bae et al. (2025) propose Mixture-of-Recursions (MoR), “a unified framework that combines [parameter sharing and adaptive computation] inside a single Recursive Transformer. MoR reuses a shared stack of layers across recursion steps to achieve parameter efficiency, while lightweight routers enable adaptive token-level thinking.” MoR offers an alternative path to per-token compute allocation that is orthogonal to memory-token augmentation.

Yu et al. (2026) (MeSH) diagnose the same fundamental issue we motivate—that the single hidden state in a recursive transformer is overloaded, forced to simultaneously carry persistent context and transient computation. They quantify this with three observables: a skewed computational pattern (the first loop performs most of the work, subsequent loops contribute negligibly), representational stagnation (high CKA similarity between consecutive loop states), and loop representational collapse (rapid singular-value decay indicating low-dimensional collapse). Their solution, Memory-as-State-Highways, “replaces the overloaded hidden state with an explicit memory buffer governed by lightweight, step-wise routers.” MeSH-enhanced recursive transformers at the 1.4B Pythia scale, with 33% fewer non-embedding parameters than the non-recursive counterpart, “improv[e] average downstream accuracy by +1.06%.” Architecturally, MeSH slots are full-sequence-shaped tensors (closer to TRM’s per-position latents than to our memory tokens), and the central architectural innovation is breaking weight-tying for the routers—each iteration uses different routing parameters—which is what enables the functional specialization across iterations they observe.

Frey et al. (2026) combine adaptive per-layer looping with gated memory banks at the ∼ 200 {\sim}200 M parameter scale on language modeling. Their architecture pairs PonderNet-style halting with two types of static learnable memory: local memory banks per layer and a global memory bank shared across layers, both retrieved via key-value attention with input-dependent gating. They report a functional dissociation: looping primarily benefits mathematical reasoning (a Loop-3 model reduces math BPB from 2.163 to 1.687, a 22% reduction), while memory banks help recover commonsense performance, and the combination outperforms an iso-FLOP baseline with three times the layers on math benchmarks. They frame this as a distinction between knowledge manipulation (improved by looping, since iterating refines representations) and knowledge capacity (improved by memory, since unique parameters store more information): “the core tradeoff is between knowledge manipulation, which looping enhances as it repeatedly refines the representations, and knowledge capacity, which requires additional unique parameters.”

The works above converge with ours from very different regimes. Yu et al. (2026) arrive at memory-based externalization as a fix for diagnosed pathologies in LM recursion; Frey et al. (2026) arrive at memory-plus-looping as a complementary architecture combining math and commonsense performance; Geiping et al. (2025) and Zhu et al. (2025) demonstrate the scaling viability of depth recurrence at LM scale; we arrive at memory tokens as the minimum architectural element required for functional behavior in a minimal UT on combinatorial reasoning. Different scales, different tasks, different memory mechanisms—the same architectural conclusion in each case: the residual stream alone is insufficient for sustained recursive computation, and some form of explicit state externalization is needed.

### 6.5 Adjacent: External Memory and Adaptive Long Context