The ability to solve harder problems by spending more time on computation is one of the most fundamental ideas in both human cognition and computer science. Autoregressive language models achieve this by writing out their reasoning in words, but the bandwidth of that channel is limited. A complementary approach is to have a distributed hidden state that gets iteratively refined, accumulating computation at every step. Looped models implement exactly this idea: they apply the same neural network repeatedly to a hidden state, growing their effective depth without growing their parameter count. In practice, however, these models have a training problem that limits how far they can go.
The training bottleneck that keeps looped models stuck
When you train a looped model, you want each update of the hidden state to produce something useful for future updates. The natural objective is backpropagation through time (BPTT), which flows gradients through every step of the recurrence. But BPTT scales linearly in memory and compute with the number of steps, and it suffers from vanishing and exploding gradients. In practice, training backpropagates through only one or two updates, which means later losses cannot provide a learning signal to earlier states. Each step must discover a globally useful computation from purely local gradients. Previous work showed this is hard in practice: looped models can learn unstable recurrences that fail even on simple problems like Sudoku, getting stuck in non-converging loops or falling into spurious attractors.
Flow and diffusion models face a related challenge from a different angle. These models learn to generate samples by denoising, starting from noise and progressively refining it. They can spend more computation during inference by taking more denoising steps, which is the same idea as looped models but applied to generation rather than reasoning. Standard training uses a single noise level per forward pass, so there is no explicit incentive for the model to build computation that transfers across steps.
Looped flows: denoising as the training signal for recurrence
Looped flows combine these two ideas in a way that solves the training problem. The core insight is that denoising objectives at different noise levels are naturally related: cleaning a heavily noisy input requires the same kind of structure recognition as cleaning a mildly noisy one, just at different scales. If you train a recurrent model on a sequence of denoising tasks with progressively decreasing noise levels and shared noise samples, the model is incentivized to learn hidden states that transfer useful computation from one step to the next, even when gradients only cover a few updates.
Here is how it works concretely. Given a problem with input c and correct solution x_1, the model receives a noisy interpolant I_t = (1-t)*x_0 + t*x_1, where x_0 is noise and t controls the noise level. A stateful denoiser takes this interpolant and a recurrent state z, predicts the clean solution, and updates the recurrent state. During training, k time steps are sampled and sorted in increasing order, so the noise levels go from high to low. Each step sees the same noise-target pair, meaning the denoising task is the same across steps but progressively easier as noise decreases. The loss sums cross-entropy predictions at each time step against the true solution.
The critical mechanism is temporal alignment. Because adjacent denoising steps share the same noise and target, the model learns that a hidden state produced while denoising at noise level t is useful input for denoising at the slightly lower noise level t+1. This gives the recurrence a natural curriculum: early steps handle the hardest denoising (most noise) and produce states that later steps can build on. The gradient does not need to flow through the entire sequence for this to work; the local alignment between adjacent objectives is enough.
An adaptive computation time (ACT) head is added on top of the recurrent state. It predicts a halting probability at each step, and during training a binary cross-entropy loss encourages the model to halt when its prediction is correct. This lets the model learn to spend different amounts of computation on different inputs.
Inference as probability flow with recurrent states
At inference time, looped flows work differently from standard looped models. Instead of just running the recurrence and taking the final prediction, the model integrates the velocity field of a probability flow. The velocity at each time step is parameterized by the learned denoiser: b_t(x; c) = (D_t(x; c) - x) / (1-t), where D_t is the denoiser's prediction. Starting from noise, the model takes ODE steps, where each step calls the denoiser with the current noisy state and the recurrent state, gets a prediction and an updated state, and uses the prediction to update the noisy state.
This formulation has two practical advantages. First, inference-time computation scales naturally: you can take a finer temporal grid with more steps to solve harder problems, and the model's accuracy improves accordingly. On Sudoku-Extreme, accuracy goes from 74.5% at 8 steps to 97.9% at 128 steps. Second, you can add stochasticity by using an SDE instead of an ODE, which injects noise through a Brownian motion term. Different initial noise samples then produce different predictions, giving you multiple valid solutions from a single model. On N-Queens 8x8, this gives 99.9% accuracy with 91.4% coverage (finding 91.4% of all valid solutions across 20 inference runs). On Graph Coloring 8 vertices, it finds solutions with only 0.7 conflicts on average and 89.4% coverage.
Architecture and training cost
The base architecture is the Trained Recurrent Model (TRM), a looped transformer with two recurrent states: one used for prediction decoding and one for internal updates. The shared network has 2 layers with hidden width 512, SwiGLU channel MLPs of width 1,536, 8-head noncausal attention with rotary position embeddings, and RMSNorm after each residual addition. Total parameter count is 5-7M depending on the task. The added parameters from the flow formulation (the noisy interpolant projection and time encoding) account for about 0.27M, or roughly 5% of the total.
Training is efficient. Sudoku, Maze, N-Queens, and Graph Coloring take 1-5 hours on a single GPU. ARC-AGI tasks take 1-2 days. The authors used institutional GPU clusters with H100, L40S, A40, A10, and RTX 3090 GPUs. Task-specific puzzle embeddings for ARC are optimized with signSGD at learning rate 10^-4. A pseudotarget regularizer (probability increasing linearly from 0 to 1 over 20k steps) helps on Sudoku and ARC-AGI-2.
Benchmark results: beating every prior looped model
Across six reasoning benchmarks, looped flows outperform prior state-of-the-art looped models on five and remain competitive on the sixth. The headline numbers: 58.8% test accuracy on ARC-AGI-1 (up from 44.6% for TRM and 47.5% for FPRM) and 12.2% on ARC-AGI-2 (up from 7.8% for TRM). On Sudoku-Extreme, 97.9% accuracy at 128 inference steps. On Maze-Hard, 86.7%, competitive with FPRM's 87.0%.
With inference-time ensembling (selecting among 5 trajectories using the ACT halting score), the numbers improve further: 59.5% on ARC-AGI-1, 99.3% on Sudoku-Extreme. On multi-solution benchmarks, looped flows achieve 99.9% accuracy on N-Queens 8x8 with 91.4% coverage, and 94.4% accuracy on N-Queens 10x10 with 61.5% coverage, both best among all methods tested including autoregressive transformers (96.3% and 90.0% respectively) and GRAM (99.7% and 89.7%).
On the Graph Coloring benchmark, looped flows find near-optimal solutions: 0.7 conflicts on 8-vertex graphs (best among all methods) and 1.0 conflicts on 10-vertex graphs, with 89.4% and 55.2% coverage respectively. The autoregressive AR transformer, despite having 10.6M parameters, gets 19.0 and 61.3 conflicts on the same tasks, showing that the recurrent approach with 7M parameters is substantially better at constraint satisfaction.
What the ablation reveals
The ablation study decomposes where the gains come from. Removing time conditioning drops ARC-AGI-1 by 2.4 points and ARC-AGI-2 by 2.3. Removing the interpolant (the noisy input formulation) drops ARC-AGI-1 by 7.3 points, the largest single component. Removing both time conditioning and interpolant together drops ARC-AGI-1 by 15.2 points and ARC-AGI-2 by 7.2, bringing performance close to the standard looped model baseline of 44.6%. Removing decreasing noise levels or noise sharing each individually have smaller effects, but they are what make the temporal alignment work.
On the inference side, switching from ODE to SDE improves accuracy on the harder tasks. On N-Queens 10x10, SDE gets 94.4% accuracy versus 73.2% for ODE, because the stochastic exploration helps escape local optima. On Graph Coloring 10 vertices, SDE reduces conflicts from 2.7 to 1.0 while also improving coverage from 54.0% to 55.2%.
Why the recurrence stays stable
One of the most striking results comes from analyzing the learned recurrence on Sudoku. The authors compared TRM and looped flows on approximately 65,000 test instances. TRM fails on 12.6% of cases: 88.3% of failures are non-convergence (the recurrence never settles on a valid solution) and 11.7% are spurious attractors (the recurrence gets stuck in an invalid fixed point). Looped flows recover 89.9% of the non-convergence cases and 98.0% of the spurious attractor cases, resolving 90.9% of TRM failures overall.
The stability comes from the training objective. Because each step of the recurrence is trained on a denoising task that is related to but easier than the previous step, the hidden state naturally evolves toward a useful representation. The denoising objective provides a dense supervision signal at every step, unlike standard looped training where the signal can be weak or absent for early updates.
Limitations and what comes next
The approach has several limitations worth noting. First, the model size is small (5-7M parameters) and the benchmarks are synthetic. Whether the ideas scale to larger models and real-world reasoning tasks remains to be shown. Second, the inference-time scaling on Sudoku (from 74.5% at 8 steps to 97.9% at 128 steps) suggests there is headroom, but the relationship between compute and accuracy is not characterized beyond this single curve. Third, the multi-solution coverage on harder instances like N-Queens 10x10 (61.5%) and Graph Coloring 10 vertices (55.2%) leaves substantial room for improvement.
The authors identify three promising directions. First, scaling to larger models and more complex reasoning tasks. Second, exploring the connection between the flow formulation and other generative approaches like energy-based models, which have also been applied to reasoning. Third, investigating whether the temporal alignment mechanism can be made adaptive rather than following a fixed noise schedule.
For practitioners, the key takeaway is that the flow formulation is not just a generative tool. It provides a natural training objective for recurrent computation, where denoising at different noise levels serves as a curriculum that teaches the hidden state to transfer information across steps. The inference-time scaling property is particularly interesting: you can trade compute for accuracy in a smooth, predictable way, and the stochastic variant gives you diverse solutions without retraining. At 5-7M parameters and hours of training, this is an approach that researchers with modest compute budgets can actually experiment with. Read the paper on arXiv