One Step Forward and K Steps Back: Better Reasoning with Denoising Recursion Models
Abstract
Looped transformers scale computational depth without increasing parameter count by repeatedly applying a shared transformer block and can be used for iterative refinement, where each loop rewrites a full fixed-size prediction in parallel. On difficult problems, such as those that require search-like computation, reaching a highly structured solution starting from noise can require long refinement trajectories. Learning such trajectories is challenging when training specifies only the target solution and provides no supervision over the intermediate refinement path. Diffusion models tackle this issue by corrupting data with varying magnitudes of noise and training the model to reverse it in a single step. However, this process misaligns training and testing behaviour. We introduce Denoising Recursion Models, a method that similarly corrupts data with noise but trains the model to reverse the corruption over multiple recursive steps. This strategy provides a tractable curriculum of intermediate states, while better aligning training with testing and incentivizing non-greedy, forward-looking generation. Through extensive experiments, we show this approach outperforms the Tiny Recursion Model (TRM) on ARC-AGI, where it recently achieved breakthrough performance.
Lay Summary
To solve difficult reasoning problems, today’s AI models usually rely on a very large “brain”, with billions of connections between artificial neurons. Recently, however, the Tiny Recursion Model (TRM) showed that strong performance on challenging puzzles, including some that are difficult for humans, is possible with only a few million connections, more than 1,000 times fewer than standard models. The key idea is recursive refinement. Instead of using a long stack of different layers, the model runs the same small network repeatedly in a loop and gradually improves its answer step by step. A major challenge for these recursive models is training. They are asked to learn hard reasoning problems from scratch, without first seeing easier versions, like being dropped into a calculus class before learning algebra. Our method combines recursion with a training approach called diffusion. During training, we give the model the correct answer with some noise added to it, so the answer is partly hidden or corrupted. When only a little noise is added, the problem is easier because much of the answer is still visible. When more noise is added, the model has to do more of the reasoning itself. We then train the model to remove this noise over multiple recursive steps, which gives it a gradual path toward learning more complex puzzle-solving behavior. We show that this approach improves on TRM and outperforms the best open-source baseline when controlling for the training data used. We also find that performance continues to improve as the model gets larger and as it is trained on more data. These results suggest that diffusion-based recursion models could benefit from scaling in the same way as large language models such as ChatGPT, opening the door to further advances in compact but powerful reasoning systems.