Symbol-Equivariant Recurrent Reasoning Models

AI summary

Overview

Research area: Machine learning / neural reasoning architectures — specifically recurrent reasoning models (RRMs) for structured, constraint-based puzzle tasks such as Sudoku, ARC-AGI and mazes.

Technical level: Advanced. The paper builds on grouped-equivariance concepts, axial attention and the RRM family (HRM, TRM), and assumes familiarity with transformer blocks, symmetry/equivariance and deep supervision.

One-sentence scope: The paper introduces SE-RRMs, a recurrent reasoning architecture that enforces permutation equivariance over symbols directly in the network layout rather than through data augmentation, and evaluates it on Sudoku, ARC-AGI and maze solving.

What This Paper Is About

Structured reasoning problems such as Sudoku and ARC-AGI involve symbols (digits, colors) that are interchangeable — permuting the symbols should permute the solution in exactly the same way — yet existing recurrent reasoning models such as the Hierarchical Reasoning Model (HRM) and Tiny Recursive Model (TRM) only capture this symmetry implicitly, by augmenting training data. The authors ask whether embedding that symbol symmetry into the architecture itself produces models that are more accurate, more parameter-efficient, and able to generalize to Sudoku grid sizes never seen during training. The proposed answer is the Symbol-Equivariant Recurrent Reasoning Model (SE-RRM), which represents the problem state as a three-dimensional tensor linking positions with symbols and applies attention across both dimensions.

Key Contributions

1. A new architecture class: The authors introduce Symbol-Equivariant Recurrent Reasoning Models (SE-RRMs), recurrent reasoning architectures that are explicitly equivariant under permutations of equivalent symbols, such as digits in Sudoku or colors in ARC-AGI problems.
2. Reduced reliance on data augmentation: They demonstrate that SE-RRMs reduce reliance on heavy data augmentation. Where symbol equivariance is not handled with extensive augmentation, SE-RRMs perform substantially better than other RRMs while using only a fraction of the trainable parameters (a 2-million-parameter SE-RRM is used throughout the experiments).
3. Architectural guarantees and inference-time flexibility: They show that symbol equivariance at the architectural level enables new symbols to be integrated during inference while preserving or improving the test-time scaling properties of recurrent reasoning.
4. Formal equivariance properties: Two propositions establish that an SE-RRM block without positional embeddings is equivariant to permutations of input positions, and that an SE-RRM block without task-type embeddings is equivariant to permutations of input symbols.

Main Findings

• Sudoku at the training size: On 9×9 Sudoku (the training regime for all RRMs), SE-RRM reaches a fully solved rate (FSR) of 93.73 (95% Wilson interval 93.66–93.81) and grid-point accuracy (GPA) of 97.58 (97.54–97.63), compared with TRM at 71.94 FSR / 89.80 GPA and HRM at 63.53 FSR / 86.11 GPA.
• Generalization to smaller and larger grids: Trained only on 9×9 puzzles, SE-RRM reaches 95.46 FSR and 99.15 GPA on 4×4 Mini Sudoku, 51.95 GPA (0 FSR) on 16×16 Maxi Sudoku, and 31.49 GPA (0 FSR) on 25×25 Ultra Sudoku. HRM and TRM are marked as unable to extrapolate to the 16×16 and 25×25 sizes at all.
• Failure of non-equivariant RRMs on 4×4: HRM and TRM both score 0 FSR on 4×4 puzzles (GPA 29.19 and 45.88 respectively), while SE-RRM solves 95.46% exactly.
• Large language model baseline: GPT-OSS-20B solves 100% of 4×4 Mini Sudoku (FSR 100, GPA 100) but only 21.67 FSR on 9×9 Sudoku (its GPA is not computed because the model sometimes refuses to produce output), and it refuses to provide solutions on 16×16 and 25×25. LLM-based systems were evaluated without task-specific program synthesis or symbolic search.
• ARC-AGI: The abstract reports competitive performance on ARC-AGI-1 and ARC-AGI-2 with substantially less data augmentation and only 2 million parameters. Specific ARC-AGI scores are not reported in the portion of the paper available here.
• Computational cost: Attention in a standard implementation costs O(I²) compute and O(I) memory; in SE-RRM this becomes O(I²K + K²I) compute and O(IK) memory. Since I ≫ K for the problems considered, cost grows linearly by a factor of K — but in domains where K ≫ I, cost is dominated by O(K²I), which the authors note may render SE-RRM impractical.
• Training protocol comparison: HRM and TRM use a Q-learning stopping policy for deep supervision; SE-RRM replaces this with a simpler stochastic mechanism that terminates training-time supervision with probability p at each deep supervision step unless the maximum number of steps has been reached. All models execute the full maximum of 16 deep supervision steps at inference.

Methodology in Plain English

The starting point is the RRM family: a transformer-style block is applied repeatedly to a recurrent state, with a fixed task embedding present at each step, and supervision applied at intermediate steps (deep supervision). In a vanilla RRM, each symbol (e.g. each Sudoku digit) gets its own learned embedding vector, and the state is a two-dimensional D × I matrix of features over positions.

SE-RRM changes this in three ways. First, the symbol embedding becomes position-independent: the same vector d is used for a given symbol everywhere, with separate vectors for special symbols such as masks and unknown tokens, and a zero embedding where a symbol is absent. Second, the state gains a third dimension, becoming a D × I × K tensor that explicitly links positions to the K possible symbols. Third, each block contains two self-attention operations instead of one: one applied along the position dimension (T^{D,I}) and one applied along the symbol dimension (T^{D,K}), each followed by normalization and a SwiGLU position-wise MLP applied token-wise. Because the symbol embedding is shared, the attention over symbols is permutation-equivariant, and because all other operations act uniformly along the symbol and position dimensions, the block preserves that equivariance — which the authors prove in Propositions 2.2 and 2.3. Task-type information (relevant for multi-task datasets such as ARC-AGI) is encoded by a learned embedding of shape D × 1 × K that is broadcast along positions and initialized identically across symbols so that equivariance holds at the start of learning. RoPE2d is used instead of RoPE to encode spatial positions. Output predictions are produced by a linear map that reduces the feature dimension to a single logit per position-symbol pair.

For the Sudoku experiments, models are trained on a subset of 1,000 samples of the 9×9 Sudoku dataset from Wang et al. (2025) with 1,000 augmentations per training sample, validated on the full test set of 422,786 samples, and then tested zero-shot on 4×4 (all 288 unique combinations), 16×16 (216 puzzles with unique solutions generated with py-sudoku) and 25×25 grids (42 samples extracted from curated puzzles in Sudoku-SMT-Solvers; the Z3Solver solved only 42 of 100 puzzles). Performance is reported as fully solved rate and grid-point accuracy over initially unfilled cells, with 95% Wilson score confidence intervals.

Why This Matters

Encoding symmetry into the architecture rather than the data changes the economics of training reasoning models: the same accuracy is reached with far fewer parameters and far less augmentation, and the model gains a form of generalization (across symbol cardinalities and, per the authors, to newly introduced symbols at inference) that augmentation cannot provide.

Real-world applications suggested by the paper's framing:

• Combinatorial optimization and scheduling, where problem sizes and structures vary widely and symbols are interchangeable.
• Planning and constraint-satisfaction tasks, which the authors describe as sharing the constraint-satisfaction and combinatorial structure central to many practical problems.
• Clinical and diagnostic reasoning scenarios, cited among the reasoning challenges where deep networks and LLMs struggle.
• Insurance pricing, where the authors note TRMs have already shown practical promise.

Industry relevance: the work argues that compact, symmetry-aware recurrent models can act as cheap, interpretable neural solvers that need neither large-scale pretraining nor reinforcement learning and where reasoning depth can be scaled at inference time independently of model capacity. Because the strongest ARC-AGI-2 results currently rely on heavily orchestrated, LLM-guided symbolic pipelines with extensive human-designed components and high computational cost, an alternative that encodes symmetry directly in a 2-million-parameter network is a meaningful counterpoint.

Future Directions

• Complexity limits when K ≫ I: The authors note that in domains where the symbol set is much larger than the number of positions, the O(K²I) symbol-attention term may render SE-RRM impractical — motivating cheaper symbol-mixing mechanisms.
• Equivariance versus task-type flexibility: Task-type embeddings are deliberately initialized identically across symbols to preserve equivariance at the start of learning but allowed to deviate afterwards; when symbol equivariance is not appropriate within a single task type remains an open design question.
• Ablations and hyperparameters: The paper defers a detailed hyperparameter comparison and several ablations (including the effect of the number of deep supervision steps) to the appendix, leaving room for systematic study of the stochastic deep-supervision termination probability p.
• Broader evaluation: The abstract mentions ARC-AGI-1, ARC-AGI-2 and mazes in the experimental setup; detailed results for these beyond the high-level ARC-AGI claim are not present in the material summarized here, and the extrapolation gaps on 16×16 and 25×25 puzzles (0 FSR despite above-chance GPA) indicate the reasoning itself, not just the symmetry, is the remaining bottleneck.

Target Audience

Researchers and practitioners working on neural reasoning architectures, structured prediction, equivariant or geometric deep learning, and efficient alternatives to large language models for constraint-satisfaction problems. It is also relevant to engineers applying recursive or recurrent reasoning models to scheduling, combinatorial optimization and pricing, and to readers tracking progress on ARC-AGI-style benchmarks who want to see how far architecture-level symmetry, rather than scale or orchestration, can go.

Authors’ abstract

Reasoning problems such as Sudoku and ARC-AGI remain challenging for neural networks. The structured problem solving architecture family of Recurrent Reasoning Models (RRMs), including Hierarchical Reasoning Model (HRM) and Tiny Recursive Model (TRM), offer a compact alternative to large language models, but currently handle symbol symmetries only implicitly via costly data augmentation. We introduce Symbol-Equivariant Recurrent Reasoning Models (SE-RRMs), which enforce permutation equivariance at the architectural level through symbol-equivariant layers, guaranteeing identical solutions under symbol or color permutations. SE-RRMs outperform prior RRMs on 9x9 Sudoku and generalize from just training on 9x9 to smaller 4x4 and larger 16x16 and 25x25 instances, to which existing RRMs cannot extrapolate. On ARC-AGI-1 and ARC-AGI-2, SE-RRMs achieve competitive performance with substantially less data augmentation and only 2 million parameters, demonstrating that explicitly encoding symmetry improves the robustness and scalability of neural reasoning. Code is available at https://github.com/ml-jku/SE-RRM.