PaperlayerView exact revision on arXiv

Source-linked AI summary

Probabilistic Tiny Recursive Model

Amin Sghaier, Ali Parviz, Alexia Jolicoeur-Martineau

arXiv:2605.19943v1cs.AI

TL;DR

Deterministic TRM inference can trap answers in bad latent-space basins, despite multiple trajectories being able to reach a unique correct answer. PTRM samples noisy parallel rollouts and selects among them with the model’s Q head. Across puzzle benchmarks, it reports accuracy gains, including nearly double the accuracy of frontier LLM ensembles on PPBench at less than 0.0001x the cost.

  • Problem

    TRM’s deterministic inference can trap incorrect answers in bad latent-space basins, although distinct latent trajectories can reach the same unique correct answer.

    Source: §1 · p. 2

  • Method

    PTRM adds Gaussian noise to latent states to create parallel trajectories and selects the candidate answer with the model’s Q head, without retraining or task-specific augmentation.

    Source: §1 · p. 2

  • Results

    91.2% vs. 55.1%: PTRM nearly doubles the accuracy of a frontier LLM ensemble on PPBench at less than 0.0001x the cost.

    Source: §7 · p. 9

  • Takeaways & Limitations

    PTRM scales TRM test-time compute through width and requires no retraining or task-specific data augmentations.

    Source: §7 · p. 9

    Experiments focus on reasoning puzzles, a subset of PPBench, and small grids; performance gains are not guaranteed across problem types.

    Source: §7 · p. 9

Abstract

from arXiv · show

Tiny Recursive Models (TRM) solve complex reasoning tasks with a fraction of the parameters of modern large language models (LLMs) by iteratively refining a latent state and final answer. While powerful, their deterministic recursion can lead to convergence at suboptimal solutions, without escape mechanism. A common workaround relies on task-specific input perturbations at test time combined with answer aggregation via voting. We introduce Probabilistic TRM (PTRM), a task-agnostic framework for test-time compute scaling that addresses this limitation through stochastic exploration. PTRM injects Gaussian noise at each deep recursion step, enabling parallel trajectories to explore diverse solution basins, and selects among them using the model's existing Q head (used for early stopping in the original TRM). Without requiring retraining or task-specific augmentations, PTRM enables substantial accuracy gains across benchmarks, including Sudoku-Extreme (87.4% to 98.75%) and on various puzzles from Pencil Puzzle Bench (62.6% to 91.2%). On the latter, PTRM achieves nearly double the accuracy of frontier LLMs (91.2% vs. 55.1%) at less than 0.0001x the cost, using only 7M parameters.

1 Introduction

TRM recursively refines latent states instead of generating token-level reasoning, but its deterministic inference can trap answers in incorrect basins. PTRM adds stochastic test-time trajectories and selects among them with the model’s Q head, without retraining or task-specific augmentations.

1 Introduction

  • TRM recursively refines a latent state rather than autoregressively generating token-level reasoning, producing one deterministic answer per input.This fits tasks with unique answers but offers no variation among trajectories for a given input.Source: §1 · p. 2
  • Many incorrect TRM answers arise from rollouts trapped in bad latent-space basins that deterministic recursion cannot escape.The paper proposes stochastic trajectories as a way to explore alternative basins and select among resulting answers.Source: §1 · p. 2
  • PTRM runs parallel rollouts with Gaussian noise at each deep recursion step and uses the Q head to select a candidate answer, without retraining or task-specific augmentation.The Q head was jointly trained as a correctness classifier but is conventionally used only for adaptive computation during training.Source: §1 · p. 2

2 Background: Tiny Recursive Model

TRM uses one network to recursively update a reasoning latent and predicted answer, with deep supervision carrying states forward across training steps. At inference, recursion can be extended for depth scaling, while the jointly trained Q head supports training-time early halting.

2 Background: Tiny Recursive Model

  • TRM iteratively refines predicted answer y to question x through recursive updates of reasoning latent z using the same two-layer network.Each latent recursion updates z n times, then updates y once.Source: §2 · p. 2
  • Deep recursion runs T latent recursions in sequence, retaining gradients only for the final one to enable large effective depth with efficient training.The network distinguishes update types by whether the input includes x.Source: §2 · p. 2
  • Deep supervision carries detached latent and answer states into the next step, training predictions to improve progressively across up to Nsup steps.At inference, the recurrence can be unrolled beyond training length for additional depth scaling.Source: §2 · p. 2
  • The Q head is trained with a correctness loss and allows the supervision loop to halt early when sigmoid(ˆq) > 0.5.During inference, the described baseline does not use the Q head and runs Nsup supervision steps.Source: §2 · pp. 2–3

3 Problem: When Does TRM Fail?

TRM trajectories include quick successes, delayed escapes from bad basins, and failures that remain trapped. Across these modes, Q tracks cell accuracy and separates correct from incorrect trajectories at convergence, motivating its use to identify better trajectories.

3.1 Analysis of failures and successes

  • Quick-success trajectories reach a convergence region within a few steps, with cell accuracy and Q rising and saturating together.They remain in that region thereafter.Source: §3.1 · p. 3
  • Delayed-success trajectories oscillate in one region for multiple steps before sharply escaping; Q and cell accuracy spike when they escape.Q is initially negative during the oscillation phase.Source: §3.1 · pp. 3–4
  • Failure trajectories oscillate within a bounded region without converging, with cell accuracy below near-perfect and Q negative throughout.Source: §3.1 · p. 4
  • Initially, delayed successes and failures both occupy bad basins with negative Q; delayed successes later escape to a good basin while failures remain stuck.Basins are latent-space regions sustained across supervision steps with similar cell accuracy.Source: §3.1 · p. 4

3.2 The Q head tracks trajectory quality

  • At convergence, Q logits separate correct from incorrect trajectories near +6 and −6, respectively, while mean Q and cell accuracy rise together on correct trajectories and stay mostly flat on incorrect ones.The aggregate pattern across 100 PPBench validation puzzles supports Q as an indicator of reaching a good basin.Source: §3.2 · p. 4, Figure 3 · p. 4
  • Because deterministic TRM cannot sample distinct trajectories, the authors propose adding Gaussian noise to sample trajectories and use Q to identify better ones.Source: §3.2 · p. 4

4 Method: Test-Time Compute Scaling via Stochastic Rollouts

PTRM makes inference stochastic by adding Gaussian noise to latent states, running K rollouts, and selecting with the Q head. Width scaling improves PPBench validation accuracy, while depth adds gains on some puzzle types and is sequential rather than parallel.

4 Method: Test-Time Compute Scaling via Stochastic Rollouts

  • PTRM adds Gaussian noise scaled by σ to latent states at each supervision step, scores candidates with Q, and decodes the highest-scoring candidate.It requires no special training or task-specific augmentation and can be applied to pretrained TRM models.Source: §4 · p. 4, Figure 4 · p. 5

4.1 Escaping bad basins

  • 8% of 100 rollouts escape a failed puzzle’s bad basin and produce correct answers, while 92% remain stuck.No rollouts escape at K=5, one escapes at K=25, and eight at K=100.Source: §4.1 · pp. 4–5

4.2 Width scaling

  • 76.4% to 89.5%: on held-out PPBench validation puzzles, pass@K and best-Q@K both rise across K=1 to K=100, with less than 1pp between them.Mode@K rises by only 1.3pp over the same range, and the measurements average three seeds.Source: §4.2 · p. 5
  • Tripling deterministic depth from 16 to 48 steps raises PPBench validation accuracy from 76.4% to 79.5% (+3.1pp), while higher K adds depth gains only on some puzzle types.At K=100, Sudoku gains +4pp; width is more practical because rollouts are parallelizable whereas added depth is sequential.Source: §4.2 · p. 5
  • PTRM provides a task-agnostic recipe for scaling TRM test-time compute across benchmarks and against baselines including frontier LLMs.Source: §4.2 · p. 5

5 Experiments

The experiments evaluate PTRM against deterministic TRM and frontier LLMs across puzzle benchmarks, examining stochastic width scaling, depth scaling, and Q-head selection.

5.1 Setup

  • Experiments use standard TRM checkpoints, held-out PPBench sets, and reported per-puzzle accuracy; PTRM’s rollout settings vary by benchmark.The benchmarks include five PPBench puzzle types, Sudoku-Extreme, Maze-Hard, and ARC-AGI 2; cost estimates convert PTRM H100 runtime using $2.50 per hour.Source: §5.1 · p. 6

5.2.1 Per-puzzle accuracy

  • 62.6% to 91.2%: PTRM raises aggregate PPBench golden-set best-Q@K, with stochastic width scaling providing most of the gain.The largest puzzle-type improvements are sudoku, from 46.7% to 97.8%, and tapa, from 40.0% to 80.0%; depth scaling alone gives a small boost.Source: §5.2.1 · pp. 6–7, Table 1 · p. 7

5.2.2 Comparison with frontier LLMs on golden set

  • 91.2% vs. 55.1%: PTRM exceeds the seven-model LLM ensemble on PPBench golden-set accuracy, despite the ensemble’s stacked advantages.The ensemble assumes a perfect verifier and counts a puzzle solved if any model solved it under any strategy.Source: §5.2.2 · p. 7, Table 2 · p. 8

5.3 Sudoku-Extreme, Maze-Hard, and ARC-AGI-2

  • 99.06% pass@K and 98.75% best-Q@K on Sudoku-Extreme; on Maze-Hard, PTRM reaches 95.63% pass@K and 85.17% best-Q@K.The deterministic Sudoku-Extreme baseline is 87.3%; on Maze-Hard it is 83.8%. The reported table covers Sudoku-Extreme, Maze-Hard, and ARC-AGI-2, with PTRM improving over standard TRM on all three.Source: §5.3 · p. 7, Table 3 · p. 8

5.4 Q head selection as σ grows

  • PTRM finds correct solutions beyond deterministic inference on every tested dataset, but the Q head selects them less reliably on Maze-Hard than on other tasks.On PPBench and Sudoku-Extreme, best-Q@K is within a point of saturated pass@K; the gap on Maze-Hard leaves headroom for a stronger verifier.Source: §5.4 · p. 8

6 Related Work

Related work situates PTRM among recursive reasoning architectures and latent-space methods, distinguishing its test-time noise mechanism from approaches requiring retraining or task-specific interventions.

6 Related Work

  • Recursive computation has been explored in Universal Transformers [10], Mixture-of-Recursions [11], Deep Thinking models [12] [14], and HRM [2].Related work also describes TRM training acceleration [15] and extensions of TRM-style architectures to language modeling.Source: §6 · p. 9
  • Latent-space methods include continuous thinking tokens derived from CoT traces [17] [18] and learning by superposition [19], which avoids explicit token sampling.These approaches may mitigate unfaithfulness and backtracking in standard autoregressive reasoning [20] [21].Source: §6 · p. 9
  • PTRM adds Gaussian noise at every supervision step, unlike generative TRM’s retraining-based sampling and concurrent work’s initial-state-only noise.Task-specific augmentation, inference-time input perturbations, and checkpoint bootstrapping address related attractors in HRM [5], whereas PTRM uses a single deterministic checkpoint across multiple puzzle types.Source: §6 · p. 9

7 Conclusion

PTRM scales test-time compute through parallel exploration and selection, without retraining or task-specific augmentation. It improves accuracy across puzzles, though its scope and gains vary across problem types.

7 Conclusion

  • PTRM scales test-time compute through width-based parallel exploration and selection, requiring no retraining or task-specific data augmentation.The conclusion reports gains from scaling both width and depth, with larger gains than depth scaling alone.Source: §7 · p. 9
  • 91.2% accuracy on PPBench costs $0.001, versus 55.1% accuracy at $38.51 for an ensemble of SOTA LLMs.PTRM also improves Sudoku from 87.4% to 98.75%, Maze-Hard from 83.80% to 86.73%, and ARC-AGI from 7.8% to 8.47% pass@1.Source: §7 · p. 9
  • Experiments focus on reasoning puzzles, a subset of PPBench, and small-grid puzzles; performance gains are not guaranteed across problem types.The authors note smaller accuracy gains on ARC-AGI-2 and Heyawake.Source: §7 · p. 9
  • The authors suspect harder-to-verify problems may benefit less because the Q head may struggle to distinguish correct from incorrect solutions; stronger verifiers are future work.Source: §7 · p. 9

A Implementation Details

Experiments use one H100 GPU, and the implementation retains the standard TRM architecture and training recipes while varying the backbone by benchmark.

A.1 Compute

  • All models are trained and evaluated on a single NVIDIA H100 80GB GPU, with PTRM adding no training cost because it runs entirely at inference.Source: §A.1 · p. 12

A.2 Models

  • Sudoku-Extreme uses the 5M-parameter TRM-MLP, while Maze-Hard, ARC-AGI-2, and PPBench use the 7M-parameter TRM-Att; layout and hyperparameters remain unchanged.Both variants follow the released standard TRM architecture and training recipes.Source: §A.2 · p. 12

A.3 PPBench dataset construction

The PPBench construction uses selected puzzle types and fixed grid sizes to create uniform-length examples, with training augmentation but unaugmented validation and golden splits.

A.3 PPBench dataset construction

  • PPBench contains 62,231 constraint-satisfaction puzzles across 94 types, with Waugh’s 300-puzzle golden set held out across 20 types.Source: §A.3 · p. 12
  • The evaluation retains six puzzle types at fixed sizes—9×9 Sudoku and 10×10 for the others—and excludes Shakashaka from per-type accuracy because deterministic TRM already reaches 100%.The retained types are Sudoku, Lightup, Nurikabe, Shakashaka, Heyawake, and Tapa; Sudoku is padded to 10×10.Source: §A.3 · p. 12, Table 4 · p. 12
  • Each training puzzle expands into 10 examples using random intermediate solve states and dihedral grid transformations; validation and golden splits are not augmented.The first training example is the unaugmented initial-state/solved pair, and the other nine are randomly sampled.Source: §A.3 · p. 12
  • The merged multi-type splits use a 294-token vocabulary and sequence length 100, with per-type counts reported in Table 4.Source: §A.3 · p. 12, Table 4 · p. 12

B Noise Ablation

Inference-noise sweeps show that rollout accuracy depends on task-specific noise levels and that Q-head selection leaves verifier headroom on some tasks. A separate Q-guided Langevin experiment found no measurable benefit from its gradient component over recurrent noise alone.

B Noise Ablation

  • Figure 7 shows pass@K increasing with inference noise up to a task-dependent level; best-Q@K tracks the ceiling closely on Sudoku-Extreme but leaves larger gaps on Maze-Hard and ARC-AGI-2.mode@K exceeds Q-head selection only on Maze-Hard; the shaded region represents accuracy a better verifier could extract.Source: Figure 7 · p. 13, B Noise Ablation · p. 13
  • On Maze-Hard, pass@K rises from 83.8% deterministically to nearly 96% by σ≈1.0, while Sudoku-Extreme is near its ceiling at σ=0.1 and ARC-AGI-2 peaks near σ=0.6.The sweep uses K=25 (K=100 for Maze-Hard), D=16, and three random seeds; Sudoku-Extreme uses a 1,000-puzzle test sample.Source: B Noise Ablation · p. 13

C Q-guided Langevin sampling

  • The Langevin experiment treated the Q-head gradient as a direction for steering latent states toward high-Q regions, motivated by visualized gradients pointing toward the good-basin region.The sampling scheme applied Langevin steps after each deep recursion step, with the number of steps N as an additional scaling axis.Source: C Q-guided Langevin sampling · p. 13, Figure 8 · p. 14
  • Making Q-guided gradients input-dependent required backpropagation through the final latent recursion; an attention-pooled Q head instead enabled head-only gradients at roughly 8× faster per step without sacrificing accuracy.Source: C Q-guided Langevin sampling · p. 14
  • At matched N, Langevin sampling and noise-only ablation had essentially identical accuracy, so the gradient added nothing measurable beyond recurrent noise.The noise-only approach is simpler, needs no retraining or test-time backpropagation, and works with any TRM checkpoint out of the box.Source: C Q-guided Langevin sampling · p. 14

D Per-puzzle accuracy on the PPBench validation set

On PPBench validation, stochastic rollouts combined with depth scaling raise aggregate best-Q@K from 76.4% to 90.4%, with the largest gains on puzzles where the deterministic baseline has more headroom.

D Per-puzzle accuracy on the PPBench validation set

Loading 2605.19943v1…