Recursive reasoning goes stochastic: 10M model scores 97% on Sudoku, every frontier LRM scores zero

Generative Recursive Reasoning

Junyeob Baek, Mingyu Jo, Minsu Kim, Mengye Ren, Yoshua Bengio, Sungjin Ahn

cs.AI

2026-05-19

GRAM makes each recursive latent transition a learned random step, trained with variational inference. A 10M model hits 97% on Sudoku-Extreme where every frontier LRM scores zero.

What problem this solves

To do long reasoning, today's neural networks mostly autoregressively unroll tokens (Chain-of-Thought). The other route is recursive reasoning: instead of emitting intermediate words, run the same transition function repeatedly over a single continuous latent state, turning "how long to think" into "how many times to loop". HRM, TRM and Looped Transformer sit on this route. They save tokens and scale with depth, but they are deterministic. One fixed latent trajectory, one answer, every time.

That breaks on two families of problems. First, problems with several valid answers (N-Queens, graph coloring, Sudoku): a deterministic model returns the same solution and its coverage is capped. Second, uncertain settings where you want to keep multiple hypotheses alive. The authors call this mode collapse, the same failure mode seen in generative models. The paper's goal is to turn recursive reasoning from a deterministic computation into a genuine probabilistic generative model.

Method

GRAM's change is small: add a learnable random perturbation to the recursive transition. The deterministic update proposes ut = fH(h, l); GRAM adds Gaussian noise ε N(μθ(ut), σ²θ(ut)I), so the final update is zt = ut + εt. The noise is not fixed-variance white noise. The mean μθ gives state-dependent steering, the variance σ²θ controls how hard to explore. Both are learned and conditioned on the current state.

Why only at the top level? GRAM's state is hierarchical, z = (h, l): the high-level h updates once per transition and carries the abstract direction; the low-level l is refined K times inside a transition to propagate constraints precisely. Randomness acts only on h, where you want to fork paths; l stays deterministic so it nails the local constraints. The ablations confirm this matters.

With randomness in, GRAM stops being a discriminator and becomes a latent-variable generative model. It models p(y|x) by marginalizing over a whole stochastic trajectory τ = (z0 → … → zT). Training is amortized variational inference on the ELBO, with an inference network qφ(τ|x,y) paired with a prior pθ(τ|x), both conditional Markov processes. For stability, gradients backprop only through the last step of each supervision step (truncated propagation), and 16 supervision points are spread along the trajectory for dense signal.

The framing hands you two inference-time scaling axes instead of one:

Swap in an empty conditioning and the same recursion models p(x) unconditionally. One architecture, both reasoning and generation.

Results

Structured reasoning (Sudoku-Extreme, ARC-AGI):

MethodParamsSudokuARC-AGI-1ARC-AGI-2
Looped TF7M61.3n/an/a
HRM27M55.040.35.0
TRM7M87.444.67.8
GRAM10M97.052.011.1
Gemini 3 Pron/an/a75.031.1

GRAM beats every deterministic recursive baseline across the board, with about a third of HRM's parameters. The frontier LRMs (o3-mini, GPT 5.2, Grok-4) all score 0 on Sudoku-Extreme; Gemini 3 Pro is strong on ARC-AGI but still zeros on Sudoku. The paper lists these LRMs as difficulty references, not fair baselines: training data, inference budget and external scaffolding do not match, so they only show how hard this benchmark is for pure language reasoning.

At matched compute (N=20 samples × 16 steps), GRAM's 97.0% still beats TRM's 90.5%; TRM needs 320 iterations to catch up. The width axis is doing real work, and it parallelizes.

Multi-solution tasks (N-Queens 8×8, graph coloring):

MethodN-Queens 8×8 acc.CoverageColoring 10-node conflicts ↓
AR96.384.861.3
MDLM96.187.212.0
HRM78.726.7n/a
TRM66.836.1n/a
GRAM99.790.33.3

Deterministic recursive baselines cap at 36.1% coverage (TRM) and fall harder as the number of valid solutions grows; that is mode collapse. GRAM holds steady on both accuracy and coverage. Graph coloring is the clearest gap: GRAM cuts conflict edges to 3.3 while AR sits at 61.3, an order of magnitude apart. The paper attributes this to recursive refinement enforcing stricter constraint propagation than sampling alone.

Unconditional generation: TRM fully mode-collapses on MNIST (FID 303.29); GRAM at 256 steps reaches IS 2.04 and FID 73.34, on par with D3PM (1000 steps, up to 55.1M params). Generating Sudoku from an empty grid, GRAM hits 99.05% validity with no external constraint checker.

Why it matters

The most valuable point is the second scaling axis. Depth scaling (more loops) is inherently serial and pays in latency; width scaling (more sampled trajectories) is inherently parallel. That means inference compute can spread horizontally instead of only stacking time. For a paradigm sold on "spend compute on demand", this is a structural addition.

The second takeaway is more general: that stochastic guidance step is plug-and-play. Ablations show it lifts Looped Transformer, HRM and TRM alike, while swapping "stochastic" for plain white noise or random init gives TRM nothing. The gain comes from the variational framework itself, not from the act of adding noise. For anyone building latent reasoning, that is a cheap improvement to try.

Third, it folds reasoning and generation into one mechanism. The same GRAM solves Sudoku and generates Sudoku; it does conditional prediction and unconditional sampling. That is a coherent bridge between latent reasoning and world models.

Finally, a contrast: a 10M model does something on Sudoku that no frontier LRM can do. This is not "small beats large"; the LLMs are still far ahead on ARC-AGI. It shows that for tightly constrained structured reasoning, unrolling language is uneconomical, and refining in a continuous latent space with stochastic sampling is far more efficient.

Limitations

The authors' top concern is training efficiency. Deep supervision runs sequentially along the trajectory, so it cannot match the training parallelism of Transformers. They admit this is the main barrier to pushing GRAM toward larger foundation models. For now this reads as "validating an architectural idea on controlled benchmarks", not something you train a big model with tomorrow.

Second is generation risk. GRAM can produce plausible-but-invalid outputs, and the paper warns that treating these as verified answers in a downstream pipeline is dangerous. In its strong suit (constraint satisfaction), a wrong solution can be costly.

Third, multi-sample inference raises compute and energy; width is not free. The experiments are all on controlled benchmarks; real or high-stakes deployment would need extra uncertainty calibration and domain safeguards, none of which is done here.

One caveat worth flagging: the LLM comparison is a difficulty reference, not a fair fight, and the authors say so. But GRAM's 11.1% on ARC-AGI-2 versus Gemini 3 Pro's 31.1% is a real absolute gap. GRAM's lead is currently confined to the contest among deterministic recursive baselines; on more open abstract reasoning it is still well behind frontier general models.

Terms

Source

What people are saying

Related papers

All paper explainers