Symbol-Equivariant Recurrent Reasoning Models

TL;DR

Symbol-Equivariant Recurrent Reasoning Models (SE-RRMs) are introduced, which enforce permutation equivariance at the architectural level through symbol-equivariant layers, guaranteeing identical solutions under symbol or color permutations.

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.

Figures (4)

• Figure 1: Comparison of a RRM, in particular TRM jolicoeur2025less, (left) and the proposed SE-RRM (right) for solving a $4\times 4$ Sudoku. For RRM, the puzzle is encoded with a symbol-specific embedding at every position, resulting in a 2D matrix. In contrast, SE-RRM introduces a third dimension to link positions and symbols and uses the same embedding for all symbols. SE-RRM is symbol-equivariant, since both Self-Attention $T^{D,I}$ and $T^{D,K}$ layers in the transformer block operate subsequently in position and symbol dimension, and the norm, addition, and MLP layers are applied token-wise.
• Figure A1: Illustration of Model Equivariance under Symbol Permutation. The model f is equivariant if applying a permutation $\pi$ to the input $X$ results in an equivalent permutation of the output, such that $f(\pi(X))=\pi(f(X))$. Here, the permutation $\pi$ swaps the symbols 1 and 4. In contrast to other reasoning models, SE-RRM guarantees permutation equivariance.
• Figure A2: Illustration of an exemplary ARC-AGI task where the colors of the squares in the input have to be exchanged. Vanilla RRMs would rely on extensive colour augmentations, while SE-RRMs can immediately learn the task type rule.
• Figure A3: Learning curves for HRM, TRM, and SE-RRM. Validation metrics HRM, TRM and SE-RRM on Sudoku, ARC-AGI-1 and Maze are shown on the y-axis. The number of samples processes during training is shown on the x-axis. SE-RRM is more data efficient since it reaches higher metrics with lower number of samples processed.

Theorems & Definitions (3)

• Proposition 2.1
• Proposition 2.2
• Proposition 2.3