Universal Transformers Need Memory: Solving Sudoku with Depth-State Trade-offs

TL;DR

Is a single Transformer block enough to solve "Extreme" Sudoku? Yes—but only if you give it a place to "think" and don't let it quit too early. This paper introduces the Universal Transformer with Memory (UTM), proving that recursive reasoning requires memory tokens as a scratchpad. By solving the "router initialization trap," the authors achieved stable training and high efficiency, reaching 57% exact-match accuracy on one of the hardest reasoning benchmarks.

The "Router Trap": Why Recursive Models Fail to Learn

The core appeal of a Universal Transformer (UT) is its ability to apply the same set of parameters over and over (recursion). To decide when to stop, we use Adaptive Computation Time (ACT).

However, the authors discovered a critical flaw: Standard initialization is a trap. With default settings, the model's "halting router" starts with a bias that makes it quit after just 2 steps. Because it never sees the benefits of going deeper, it gets stuck in a "shallow equilibrium" and never learns to solve the puzzle.

The Fix: Deep-Start Initialization. By setting the initial bias to -3, the model is forced to process the maximum number of steps (e.g., 18) at the start of training. It only learns to stop early once it has mastered the task.

Methodology: The UTM Architecture

The UTM uses a single weight-shared block (Self-Attention + SwiGLU) but adds a crucial component: Learned Memory Tokens.

As shown in the architecture, sequence tokens and memory tokens are concatenated. The memory tokens don't correspond to any part of the puzzle; they serve as a computational scratchpad where the model can store intermediate constraints and "thoughts" across iterations.

The Necessity of Memory

One of the most striking findings is the Memory-Token Curve. Without memory tokens (T=0), the model fails completely, regardless of how many steps it pouders.

• Threshold: To solve an 81-cell Sudoku, the model needs at least 8 memory tokens.
• Plateau: Between 8 and 32 tokens, performance is stable (~57% EM).
• Dilution: At 64 tokens, performance drops because the attention mechanism becomes too "diluted" to focus on the puzzle's constraints.

How the Model "Thinks": Attention Analysis

By looking under the hood at the attention maps, the authors found that heads specialize into three distinct roles:

1. Memory Readers: Sequence tokens query the memory to retrieve constraints.
2. Memory Writers: Memory tokens broadcast global state back to the sequence.
3. Constraint Propagators: Pure sequence-to-sequence heads that develop "block-diagonal" patterns corresponding precisely to Sudoku's row, column, and box rules.

As recursion progresses (Step 0 to 17), the attention structure becomes increasingly refined, mirroring the way a human solver gradually fills in a difficult grid.

Efficiency through Lambda Warmup

Training a model to be deep is great for accuracy, but expensive. The authors used Lambda Warmup—delaying the penalty for "pondering" until the model has learned the task. This allowed the UTM to match full-depth accuracy while saving 34% in compute costs.

Interestingly, these models can "over-ponder." Running the model for 36 steps at inference (even if trained on only 18) can boost accuracy by up to 14 percentage points.

Conclusion and Future Work

The UTM proves that for recursive reasoning, State (Memory) is as important as Depth. The "Deep-Start" initialization is a simple, one-line-of-code fix that could rescue many previously failed recursive Transformer projects.

While the model currently struggles with out-of-distribution generalization (memorizing the 1k puzzle protocol rather than learning the "rule of Sudoku"), its performance on the massive Sudoku-Extreme dataset sets a new benchmark for what's possible with a parameter-efficient, single-block architecture.