Symbol-Equivariant Recurrent Reasoning Models

2026-03-02Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors improved neural networks that solve puzzles like Sudoku by making the models understand that swapping symbols or colors shouldn't change the solution. They created a new type of model called SE-RRMs which handle symbol changes directly in their design instead of relying on extra training data. This new approach helps the models solve bigger and smaller puzzles than they were trained on and works well on a different reasoning task called ARC-AGI with fewer parameters and less data. Their work shows that building in this symmetry helps these reasoning models work better and more efficiently.

SudokuARC-AGIRecurrent Reasoning ModelsPermutation EquivarianceSymbol SymmetryData AugmentationHierarchical Reasoning ModelNeural NetworksGeneralizationModel Scalability
Authors
Richard Freinschlag, Timo Bertram, Erich Kobler, Andreas Mayr, Günter Klambauer
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.