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

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

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

###### 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 T{=}0 always fails, T = 8 T{=}8 reliably succeeds) followed by a stable plateau ( T = 8 T{=}8 – 32 32 , 57.4 % ± 0.7 % 57.4\%\pm 0.7\% exact-match) and a dilution boundary at T = 64 T{=}64 . Under halt-side pressure ( λ \lambda warmup), mean halt drops monotonically with memory size across the plateau (from 11.6 at T = 8 T{=}8 to 8.3 at T = 64 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 -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 \pm 0.7 vs ± 9.3 \pm 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 .

## 1 Introduction

Universal Transformers ( Dehghani et al., 2019 ) apply a single transformer block iteratively, with Adaptive Computation Time ( Graves, 2016 ) determining per-token processing depth. While theoretically appealing—arbitrary-depth reasoning with finite parameters—practical implementations have shown mixed results ( Csordás et al., 2021 ) .

We investigate the role of learned memory tokens ( Burtsev et al., 2020 ) in enabling recursive reasoning within this architecture, which we call UTM (Universal Transformer with Memory), using Sudoku-Extreme as a testbed. The title echoes Darcet et al. (2024) , “Vision Transformers Need Registers”—we present analogous evidence that our single-block UT needs memory tokens, based on extensive empirical failure without them.

Our contributions: 1. Memory-token necessity, threshold, and depth-state substitution (§ 4 ): In our single-block UT with ACT, no configuration without memory tokens achieves non-trivial performance across any tested seed, initialization, or depth mode. The optimal count shows a sharp threshold between T = 0 T{=}0 (always fails) and T = 8 T{=}8 (always succeeds) for 81-cell Sudoku, with a stable plateau through T = 32 T{=}32 . Memory tokens and ponder depth are substitutable resources at fixed accuracy: at λ = 0 \lambda{=}0 , minimum halt steps decrease monotonically with T T ( 17.7 → 16.4 → 15.5 17.7\to 16.4\to 15.5 ); under halt-side pressure ( λ = 0.001 \lambda{=}0.001 +warmup), the trade-off becomes a clean monotonic curve, with mean halt dropping from 11.60 11.60 at T = 8 T{=}8 to 10.29 10.29 at T = 32 T{=}32 to 8.25 8.25 at T = 64 T{=}64 at near-constant accuracy on the plateau ( ∼ 57 % \sim 57\% ) and a 2pp drop at the dilution boundary. We note that other recursive architectures (TRM, HRM) solve similar tasks via different mechanisms—the necessity is architecture-specific.

2. Router initialization trap (§ 3 ): We identify that both default initialization (bias = 0 =0 , p ≈ 0.5 p\approx 0.5 ) and Graves’ recommended positive bias (bias = 1 =1 , p ≈ 0.73 p\approx 0.73 ) create shallow-halt traps. In our setting, > > 70% of runs fail to escape. We propose deep-start initialization (bias = − 3 =-3 , p ≈ 0.05 p\approx 0.05 ), which inverts the assumption and resolves the issue.

3. ACT provides reliability and efficiency (§ 5 ): Fixed-depth processing with memory tokens achieves 53.4 % ± 9.3 % 53.4\%\pm 9.3\% EM (3 seeds)—high variance. ACT-enabled runs are more consistent ( 56.9 % ± 0.7 % 56.9\%\pm 0.7\% ). Lambda warmup achieves 57.0 % ± 1.1 % 57.0\%\pm 1.1\% using 34% fewer ponder steps—matching quality with significant compute savings.

4. Diagnostic framework (§ 3 ): Per-step router probability, step-weight distribution, and attention-mass logging that reveals head specialization and computation dynamics across recursive depth.

## 2 Architecture

### 2.1 Universal Transformer with Memory Tokens

Our model applies a single UniversalTransformerBlock (pre-norm attention + SwiGLU FFN) iteratively for up to K = 18 K{=}18 steps (Figure 1 ). The full sequence at each step is [ mem 1 , … , mem N , seq 1 , … , seq L ] [\text{mem}_{1},\ldots,\text{mem}_{N},\text{seq}_{1},\ldots,\text{seq}_{L}] with bidirectional attention.

Components: token + type embedding (memory vs sequence) + per-step learned positional embedding; multi-head attention with RoPE and QK-normalization; SwiGLU FFN with 8 / 3 × 8/3\times expansion; DerfNorm (our term for the normalization-free design using Derf ( Chen et al., 2025 ) ), where standard normalization layers are replaced by erf ​ ( α ⋅ x + s ) \text{erf}(\alpha\cdot x+s) with learned per-feature α , s \alpha,s ; N N learned memory vectors with indexed RoPE positions (numbered registers). Parameters: 3.2M (hidden=512, heads=8, head_dim=64, vocab=11).

### 2.2 Adaptive Computation Time

Following Graves (2016) , each token maintains cumulative halting probability. At step k k , the router (linear + sigmoid) outputs p k ∈ ( 0 , 1 ) p_{k}\in(0,1) . The output is a weighted blend: output = ∑ k = 1 N w k ⋅ h k \text{output}=\sum_{k=1}^{N}w_{k}\cdot h_{k} , where w k = p k w_{k}=p_{k} for intermediate steps and w N = 1 − ∑ k < N p k w_{N}=1-\sum_{k<N}p_{k} at halt. Ponder cost ρ = N + R \rho=N+R is minimized with coefficient λ \lambda .

When ACT is disabled, the model outputs only the final representation h K h_{K} (standard weight-tied transformer).

### 2.3 Deep-Start Initialization

With default framework initialization (zero bias), the router computes σ ⁡ ( W ⋅ h + 0 ) ≈ 0.5 \sigma(W\cdot h+0)\approx 0.5 , causing tokens to halt after ∼ 2 {\sim}2 steps. Graves (2016) recommends initializing the halting bias to a positive value ( b h = 1 b_{h}=1 , giving σ ⁡ ( ⋅ ) ≈ 0.73 \sigma(\cdot)\approx 0.73 ) to “prevent very long sequences at the beginning of training”—which makes tokens halt even faster ( ∼ 1 {\sim}1 – 2 2 steps). Dehghani et al. (2019) does not specify halting initialization (Appendix C shows architecture but omits init details).

Both the default (bias = 0 =0 , p ≈ 0.5 p\approx 0.5 ) and Graves’ recommendation (bias = 1 =1 , p ≈ 0.73 p\approx 0.73 ) produce shallow halting. We propose the opposite:

Deep start: bias = − 3 =-3 , giving σ ⁡ ( W ⋅ h − 3 ) ≈ 0.05 \sigma(W\cdot h-3)\approx 0.05 . Tokens process all K K steps by default and learn to halt earlier. This inverts Graves’ assumption: instead of preventing long sequences, we start with maximum depth and let the model discover where to stop. This is appropriate when the task requires significant depth—the cost of long initial sequences is small if the final learned policy uses them, whereas a shallow starting policy cannot easily discover that depth is needed.

## 3 The Router Initialization Trap

### 3.1 Diagnosis

We instrumented the ACT loop to log per-step router probability and router-specific gradient norm. Across 13 completed runs (planned 5 memory-token counts × \times 3 seeds; only T = 64 T{=}64 / S = 0 S{=}0 was run for T = 64 T{=}64 before we moved on; bias=0, λ = 0 \lambda{=}0 , 4 epochs each):

Diagnostic findings: (1) All runs start at p ≈ 0.48 p\approx 0.48 – 0.52 0.52 , halt = 2.0 =2.0 . (2) By step 3k, all develop a shallow-halt pattern (halt ≈ 5 \approx 5 – 7 7 ). (3) Escape correlates with a 10–45 × \times spike in router gradient norm. (4) Stuck runs maintain router gradient < 0.04 <0.04 for all 60k steps.

### 3.2 Deep Start Resolves the Trap

With bias = − 3 =-3 , all previously-failing configurations at seed 123 succeed (Table 2 ). Training dynamics change qualitatively: with bias=0, successful runs exhibit abrupt phase transitions; with bias = − 3 =-3 , accuracy rises smoothly.

### 3.3 DerfNorm Is Not the Cause

We verify the trap is not a DerfNorm artifact by replacing it with RMSNorm at the same seed and configuration ( T = 16 T{=}16 , seed 42, bias=0): DerfNorm achieves 50.6% EM (escapes); RMSNorm achieves 0.0% EM (halt stuck at 3.5, never escapes). The trap is worse with standard normalization—it is inherent to ACT with standard initialization in our setting. This is consistent with Chen et al. (2025) ’s general finding that Derf outperforms RMSNorm across domains, though the ACT-escape dynamics we observe are specific to recursive architectures.

## 4 Memory Tokens for Recursive Reasoning

### 4.1 Task and Setup

Sudoku-Extreme ( Sapient Intelligence, 2025 ) : 9 × \times 9 puzzles, extreme difficulty (17–24 givens). 3.83M train / 423K test. Encoder-style prediction of all 81 cells simultaneously. All results: hidden=512, heads=8, max_ponder=18, batch=256, AdamW (lr= 3 × 10 − 4 3\times 10^{-4} , cosine decay), EMA (0.999), 4 epochs, bias = − 3 =-3 , λ = 0 \lambda{=}0 unless noted.

### 4.2 Memory-Token Curve

Key findings: (1) Sharp threshold between T = 0 T{=}0 and T = 8 T{=}8 : T = 0 T{=}0 always fails ( 2.5 % 2.5\% , 3 seeds). T = 4 T{=}4 is borderline—succeeds at 2 of 3 seeds (55.5–55.8%) but fails at the third (5.4%), exhibiting the same seed sensitivity as T = 64 T{=}64 . T = 8 T{=}8 always succeeds ( 57.4 ± 0.7 % 57.4\pm 0.7\% , 3 seeds). The minimum reliable scratchpad for 81-cell Sudoku is 8 tokens ( ∼ 1 {\sim}1 per 10 cells). (2) Stable plateau : T = 8 T{=}8 : 57.4 ± 0.7 57.4{\pm}0.7 %, T = 16 T{=}16 : 56.9 ± 0.7 56.9{\pm}0.7 %, T = 32 T{=}32 : 56.4 ± 1.7 56.4{\pm}1.7 %. (3) T = 0 T{=}0 fails with deep start : memory tokens are necessary in this architecture. We note that TRM ( Jolicoeur-Martineau, 2025 ) and HRM ( Wang et al., 2025 ) solve Sudoku without memory tokens via autoregressive answer improvement—the necessity is architecture-specific to our single-block UT.

### 4.3 Memory-Depth Tradeoff

Within the plateau, memory tokens and ponder depth function as substitutable resources: as T T increases, the model can solve the task at the same accuracy with fewer ponder steps. The clearest evidence is the per-seed minimum halt step in Table 3 : T = 8 T{=}8 saturates the 18-step ceiling at all 3 seeds (halt ∈ [ 17.7 , 18.0 ] \in[17.7,18.0] ), T = 16 T{=}16 admits some headroom ( [ 16.4 , 18.0 ] [16.4,18.0] ), and T = 32 T{=}32 admits more ( [ 15.5 , 18.0 ] [15.5,18.0] ). The substitution is real but seed-dependent — at T = 32 T{=}32 , S = 0 S{=}0 exploits the headroom (halt = 15.5 =15.5 ) while S = 42 S{=}42 stays saturated (halt = 18.0 =18.0 ), and different seeds find different equilibria along the same accuracy plateau.

Mean halt is dominated by the depth ceiling at λ = 0 \lambda{=}0 . Across the plateau, mean halt is 17.87 17.87 , 17.07 17.07 , 17.03 17.03 for T = 8 , 16 , 32 T=8,16,32 respectively (Table 3 ). The compression is small in the mean because λ = 0 \lambda{=}0 in this sweep — with no penalty for using deep ponder, the model saturates the ceiling whenever it can, and the substitution surfaces only in the per-seed minima. The trade-off is more visible under halt-side pressure: with λ = 0.001 \lambda{=}0.001 and a 20k-step warmup at T = 16 T{=}16 across 3 seeds (§ 5 ), mean halt drops from 17.07 17.07 to ∼ 11.4 \sim 11.4 at unchanged accuracy ( 57.0 ± 1.1 % 57.0\pm 1.1\% vs 56.9 ± 0.7 % 56.9\pm 0.7\% ). This is the canonical depth-state trade-off form: fixed accuracy, reduced compute, achieved by giving the router something to push against.

The full curve under halt-side pressure. A direct T × λ T\times\lambda +warmup sweep at S = 0 S{=}0 confirms the trade-off in its canonical form (Table 4 , Figure 4 ): mean halt decreases monotonically from 11.60 11.60 at T = 8 T{=}8 to 11.50 11.50 at T = 16 T{=}16 to 10.29 10.29 at T = 32 T{=}32 to 8.25 8.25 at T = 64 T{=}64 , with EM essentially constant across the plateau ( 57.09 → 58.00 → 56.87 57.09\to 58.00\to 56.87 ) and a 2pp drop at the dilution boundary ( T = 64 T{=}64 , 54.91 % 54.91\% ). Compared to λ = 0 \lambda{=}0 (Table 3 , S = 0 S{=}0 column), halt is compressed by 33–40% across the plateau at unchanged or slightly improved accuracy.

Total compute is approximately conserved, not strictly fungible. Total token-step compute (memory plus sequence tokens times mean halt) grows mildly with T T across the plateau ( 1032 → 1116 → 1163 1032\to 1116\to 1163 token-steps for T = 8 , 16 , 32 T=8,16,32 , ∼ 13 % {\sim}13\% growth), reaching 1196 1196 at T = 64 T{=}64 . Memory and depth substitute in halt count and per-step bandwidth, but total operations are not strictly preserved.

The dilution boundary is compatible with halt-side pressure. At T = 64 T{=}64 , halt drops further (8.25) and accuracy is 2pp below the plateau ( 54.91 % 54.91\% vs ∼ 57 % {\sim}57\% ). With only one seed in this sweep, we cannot rule out single-seed noise from the 2pp gap alone (the eval trajectory in Figure 4 b shows the four runs intermingled across the plateau). However, the magnitude and direction are consistent with the bimodal collapse observed at T = 64 T{=}64 under λ = 0 \lambda{=}0 across multiple seeds (Table 3 ); read together, they support treating T = 64 T{=}64 as a dilution boundary — an architectural property at large T T — rather than an artifact of saturated halting at λ = 0 \lambda{=}0 .

### 4.4 Fixed-Depth vs ACT Processing

Fixed-depth processing achieves comparable mean EM ( 53.4 % 53.4\% ) but with much higher variance : EM ranges from 44.9% to 63.4% across 3 seeds ( ± 9.3 % \pm 9.3\% ), compared to 56.5–57.7% for ACT ( ± 0.7 % \pm 0.7\% )—an order-of-magnitude difference in seed sensitivity. We attribute this to the output mechanism: ACT blends representations across all K K steps ( ∑ w k ​ h k \sum w_{k}h_{k} ), averaging out seed-dependent variation in individual steps; fixed-depth relies entirely on the final step’s representation h K h_{K} , which may be more sensitive to initialization-dependent optimization trajectories since it is a single snapshot rather than an average over steps. Lambda warmup combines ACT’s reliability ( 57.0 ± 1.1 % 57.0\pm 1.1\% ) with 34% compute savings.

### 4.5 How Memory Tokens Function: Attention Analysis

Attention evolves across depth: at step 0, attention is diffuse; by step 9, S → \to S develops block-diagonal patterns matching Sudoku row/column/box constraints; by step 17, the structure is highly refined (Figure 5 ).

Heads specialize into distinct roles (Figure 6 , Table 6 ). The full attention quadrant breakdown reveals asymmetric memory usage:

Three functional groups emerge: Memory readers (H4, H5: high s → \to m, sequence queries memory), memory writers (H1, H3: high m → \to s, memory broadcasts to sequence), and constraint propagators (H2, H6: s → \to m ≈ 0 \approx 0 , structured S → \to S patterns). H0 and H2 show high m → \to m (memory self-attention, ∼ 0.73 {\sim}0.73 ), suggesting internal memory coordination. H6 is striking: both s → \to s and m → \to s are ∼ 1.0 {\sim}1.0 —it propagates puzzle constraints while memory tokens passively observe the sequence.

In the trapped model (0% EM), no head develops s → \to m > 0.16 >0.16 , attention is disorganized, and the constraint patterns seen in H2/H6 are absent.

### 4.6 Small-Dataset Generalization

Following the TRM protocol (1000 training puzzles, 1000 augmentations, full 423K test set):