Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

AI summary

Overview

Research area: Efficient reasoning architectures — specifically latent recursive reasoning models and state space models (SSMs) as alternatives to attention.

Technical level: Advanced. The paper assumes familiarity with Transformer blocks, attention, Mamba/SSM state-space recurrences, normalisation placement, and the pass@K evaluation protocol.

Scope: The paper tests whether replacing the Transformer blocks inside the Tiny Recursive Model (TRM) with a parameter-matched Mamba-2 hybrid operator preserves or improves reasoning performance on ARC-AGI-1, Sudoku-Extreme, and Maze-30×30-Hard.

What This Paper Is About

Recursive reasoning models such as TRM show that very small networks (7M parameters) can solve abstract reasoning puzzles by repeatedly refining a hidden latent state instead of writing out intermediate reasoning tokens. TRM's per-step operator is an attention-heavy Transformer block, and it was not known whether a different operator — one whose own recurrence resembles iterative refinement, such as Mamba-2's state space update — could be dropped into the same recursive scaffold without losing reasoning ability. The authors swap in Mamba-2 hybrid blocks at matched parameter counts and measure whether reasoning is preserved.

Key Contributions

1. First Mamba-hybrid model for recursive latent reasoning — the paper states this is the first application of a Mamba hybrid operator inside a recursive latent reasoning architecture.
2. Empirical validation on ARC-AGI-1 — improved pass@2 performance (+2.0%) over the attention-based TRM baseline, with supplementary competitive results on Sudoku-Extreme and Maze-30×30-Hard.
3. A coverage-vs-selection analysis — pass@K patterns are decomposed into candidate coverage (does the correct answer appear anywhere in the prediction set?) and winner selection (is it ranked first?), with supporting vote statistics and a difficulty-stratified breakdown.
4. A stated argument for post-norm in recursion — the authors argue post-norm is needed for recursion stability itself, independently of the Q-learning objective that originally motivated it in HRM, because it bounds the residual stream magnitude across repeated unrolling.

Main Findings

• Pass@2 improves by +2.0%: On ARC-AGI-1, the hybrid (TR-mamba2attn, 6.86M parameters) reaches 45.88% pass@2 versus 43.88% for the attention baseline (TRM-attn, 6.83M parameters). Pass@2 is described as the official ARC-AGI metric.

• The advantage grows with K: Across the reported pass@K table, the hybrid leads at K=5 (51.88% vs 49.25%, +2.63), K=10 (54.50% vs 52.13%, +2.37), K=100 (65.25% vs 60.50%, +4.75), and K=1000 (69.75% vs 65.50%, +4.25).

• Pass@1 is essentially unchanged: 40.50% for the hybrid versus 40.75% for TRM-attn (a −0.25 difference), indicating similar top-1 selection quality.

• Candidate diversity is higher for the hybrid: From final-checkpoint statistics, the hybrid generates 339.5 unique candidates per puzzle versus 266.6 for TRM-attn (+27%), with higher vote entropy (5.39 vs 4.56).

• The baseline is more decisive: TRM-attn concentrates 41.1% of votes on its top-1 candidate (vs 32.9% for the hybrid) and has a larger top-1 margin (32.3% vs 24.0%).

• Difficulty stratification splits the effect: Using a model-agnostic 15% correct-vote-share threshold over all 419 test inputs, 246 are classified hard (<15%) and 173 easy (≥15%). On hard inputs the hybrid gains +4.9 pp at pass@5; on easy inputs TRM-attn gains +4.6 pp at pass@1. At pass@5 the models solve partially disjoint sets (31 hybrid-only vs 23 TRM-attn-only).

• Sudoku favours dense mixing: TRM-mlp-t is best at 87.4% (5.00M parameters), followed by TR-mamba2mlpt at 84.2% (6.28M). Attention-based variants trail: TRM-attn at 72.2% and TR-mamba2attn at 66.5%.

• Maze shows the opposite pattern: Both MLP-t variants score 0.0% accuracy on the 30×30 grids (TRM-mlp-t at 19.0M and TR-mamba2mlpt at 13.24M parameters), while TR-mamba2attn reaches 80.6% versus 60.8% for TRM-attn. The authors flag these results as preliminary because training varies between 6% and 85% accuracy across checkpoints.

• A reported non-attention ARC baseline: The paper cites TRM-mlp-t at 29.6% pass@2 with 19M parameters from the original TRM paper, which the authors did not reproduce. Their TR-mamba2mlpt reaches 32.125% pass@2 with 13.24M parameters, a +2.5% improvement over that reported baseline.

• Training curves are consistent: The hybrid's pass@2 and pass@100 curves pull ahead of attention after the first few epochs and stay ahead, so the effect is described as a stable property rather than a late-training phenomenon.

Methodology in Plain English

The researchers kept TRM's recursive structure exactly as it was — three outer cycles and four to six inner loops, the same two latent states, and the same output heads — and changed only the block that performs each refinement step.

The original TRM uses attention-and-MLP Transformer blocks. The authors built two replacements. The first, TR-mamba2attn, runs a pipeline of two Mamba-2 blocks, then attention, then an MLP, so sequential state-space processing is combined with cross-position attention. The second, TR-mamba2mlpt, is the same but swaps attention for an "MLP-t" block that operates on the transposed sequence dimension to give all-to-all communication without attention. A pure Mamba model was avoided because Mamba processes information causally in one direction, and the grid-based tasks need bidirectional information flow.

To make the comparison fair, parameters were matched: hidden size 512, TRM-attn at 6.83M parameters, TR-mamba2attn at 6.86M, with Mamba-2 configured as d_state = 128, headdim = 64, and expand = 2. The authors used post-norm (RMSNorm) throughout, arguing that in unrolled recursion pre-norm lets the residual stream grow and can eventually produce NaN failures, whereas post-norm re-normalises after every residual add and bounds the hidden state regardless of recursion depth.

Evaluation followed the TRM protocol. ARC-AGI-1 inputs are expanded into roughly 880 augmentations via dihedral transformations and colour permutations, predictions are inverse-transformed and aggregated by vote count and average confidence, and pass@K is reported for K in {1, 2, 5, 10, 100, 1000}. The ARC test set contains 400 puzzles with 419 test inputs (19 puzzles have two test inputs each), producing 368,150 total test instances. Sudoku-Extreme and Maze use exact accuracy, with Sudoku on 9×9 grids and Maze on 30×30 grids.

Why This Matters

Research impact: The work widens the design space for recursive reasoning. It shows that the operator inside a recursive scaffold is not fixed to attention — an SSM hybrid can sit there without hurting reasoning, and on coverage metrics it helps. It also separates two things that pass@K usually conflates: whether the model can find the right answer across its candidate pool, and whether its voting scheme ranks that answer first. The difficulty-stratified result and the 31-vs-23 disjoint-solve split suggest the two operators are complementary rather than strictly ranked.

Real-world applications:

• Abstract reasoning benchmarks and evaluation suites that use candidate ranking, where better coverage at high K could raise measured ceilings.
• Constraint satisfaction and structured puzzle solving, where the paper's Sudoku results show dense all-to-all mixing outperforming both attention and Mamba hybrids.
• Spatial path-finding and grid-based planning, which the Maze results show is sensitive to operator choice — attention hybrids worked while dense MLP-t variants scored 0.0%.
• Efficient inference-time reasoning, since SSM operators carry linear complexity and Mamba-2 is reported to train 2–8× faster through Structured State Space Duality, which matters for architectures that trade parameters for thinking time.

Industry relevance: Both authors are affiliated with Intercom (Dublin, Ireland). The direction is relevant to organisations that need strong reasoning from small, cheap-to-serve models rather than large hosted APIs — the paper notes that TRM's 7M-parameter result outperforms many commercial LLM APIs on ARC-AGI-1. An operator that keeps reasoning quality while offering more efficient inference is directly useful for latency- and cost-sensitive deployment.

Future Directions

1. Internalise the recursion into the SSM state. The authors propose investigating whether the outer recursive loop can be absorbed into Mamba's inherent inner recurrence, rather than relying solely on outer-loop iteration.
2. Purpose-designed mixing strategies. The paper frames itself as a first step toward understanding the best mixing strategies for recursive reasoning, and the Sudoku-versus-Maze split — dense mixing wins on small 9×9 grids but fails completely on 30×30 grids — makes the choice of mixer context-dependent and unresolved.
3. Compute-normalised evaluation. The authors explicitly call for it, since parameter matching does not equalise the compute or wall-clock cost of the different operators.
4. Resolving the Maze instability. The high variance across checkpoints (6–85% accuracy) means the 80.6% Maze figure is preliminary and needs confirmation.

Target Audience

Researchers and engineers working on efficient reasoning architectures — particularly those interested in latent/recursive reasoning, looped Transformers, or state space models as attention replacements. It will also suit practitioners deciding which mixing operator to use for small-model reasoning deployment, and readers of the HRM/TRM line of work who want to know how much of TRM's behaviour depends on attention. A solid grounding in Transformer internals, SSM recurrences, and pass@K evaluation is needed to follow the methodological sections; the high-level results and the coverage-vs-selection framing are accessible without it.

Authors’ abstract

Recent work on recursive reasoning models like TRM demonstrates that tiny networks (7M parameters) can achieve strong performance on abstract reasoning tasks through latent recursion -- iterative refinement in hidden representation space without emitting intermediate tokens. This raises a natural question about operator choice: Mamba-2's state space recurrence is itself a form of iterative refinement, making it a natural candidate for recursive reasoning -- but does introducing Mamba-2 into the recursive scaffold preserve reasoning capability? We investigate this by replacing the Transformer blocks in TRM with Mamba-2 hybrid operators while maintaining parameter parity (6.83M vs 6.86M parameters). On ARC-AGI-1, we find that the hybrid improves pass@2 (the official metric) by +2.0\% (45.88\% vs 43.88\%) and consistently outperforms at higher K values (+4.75\% at pass@100), whilst maintaining pass@1 parity. This suggests improved candidate coverage -- the model generates correct solutions more reliably -- with similar top-1 selection. Our results validate that Mamba-2 hybrid operators preserve reasoning capability within the recursive scaffold, establishing SSM-based operators as viable candidates in the recursive operator design space and taking a first step towards understanding the best mixing strategies for recursive reasoning.