Thinking with Looped Flows
Ayhan Suleymanzade, Chanhyuk Lee, Floor Eijkelboom, Nicholas M. Boffi, İsmail İlkan Ceylan, Jinwoo Kim
cs.LG, cs.AI
2026-09-11
Looped Flows train a stateful denoiser with local objectives at decreasing noise. Matching TRM's architecture, they reach 58.8% on ARC-AGI-1 and 12.2% on ARC-AGI-2.
Looped models reuse the same weights to update a hidden state, stretching effective depth at inference. HRM, TRM, and FPRM already get usable scores on Sudoku, mazes, and ARC-AGI with a few million parameters. Training is the bottleneck. Full backpropagation through time grows memory with the number of steps and tends to explode or vanish. In practice each step is stop-gradient, so later losses never reach earlier states. Early updates never learn to set up later computation. The recurrence then fails to converge, or it converges to a wrong attractor.
Flow and diffusion models split "noise to answer" into a sequence of local denoising tasks. Each step has its own loss, so there is no need to backprop along the whole trajectory. That trajectory is weak at some serial computations. Looped Flows join the two ideas: local denoising becomes a curriculum that induces recurrent states worth carrying forward.
The denoiser is stateful. At each step it predicts a clean solution from the current noisy sample and updates a recurrent state z. Inference then recovers the probability-flow velocity from that prediction and takes an ODE or SDE step.
Training draws k+1 times from [0, 1] and sorts them so noise (1-t) falls over the rollout. The same problem, the same noise sample, and the same solution are shared across those k steps. Each step is cross-entropy against the clean answer, with stop-gradient between steps. Shared noise aligns adjacent denoising tasks, giving the recurrent state a reason to keep features that later steps can reuse. k is 16, matching the maximum training horizon used by TRM.
As noise falls, the loss tends to drop monotonically and later steps can saturate. An ACT head (adaptive computation time) classifies whether the current prediction already matches the answer and halts if the score exceeds 1/2, so extra steps do not overfit a solved example.
Inference can integrate on a finer grid than training; extra steps are extra compute. Independent trajectories from different initial noise can be ranked by the ACT head (best-Q). An SDE integrator with γ>0 is slightly more stable than plain Euler.
The architecture is TRM with two extras: a projection for the noisy interpolant and a time embedding, about 0.27M extra parameters. Sudoku uses a 5M MLP-Mixer; other tasks use a 7M noncausal Transformer. Internally the denoiser still carries TRM's pair of states (h, ℓ), with predictions decoded from h. One denoiser call repeats a TRM-style cycle three times, backpropagating only through the last cycle (inner repeats: 6 on Sudoku, 4 elsewhere).
Six benchmarks, four single-solution and two multi-solution. Single-trajectory numbers:
| Method | Sudoku-Extreme | Maze-Hard | ARC-AGI-1 | ARC-AGI-2 |
| TRM (same architecture) | 87.4 | 85.3 | 44.6 | 7.8 |
| FPRM | 94.2 | 87.0 | 47.5 | 6.2 |
| GRAM | - | - | 52.0 | 11.1 |
| Looped Flows | 97.9 | 86.7 | 58.8 | 12.2 |
On Maze they sit 0.3 points behind FPRM. The ARC gaps versus TRM are +14.2 and +4.4. Best-Q over 5 trajectories lifts Sudoku to 99.3 and ARC-AGI-1 to 59.5; EqR needs 128 trajectories to reach 99.8 on Sudoku. ARC-AGI-2 stays at 12.2; only the spread tightens from ±1.9 to ±0.9.
On Sudoku, 8 integration steps score 74.5% and 128 steps score 97.9%, overtaking GRAM at 32 steps. Across about 65,000 Sudoku test puzzles, TRM fails 12.6% of the time: 88.3% of those failures never converge, 11.7% land in a spurious attractor. Looped Flows recover 89.9% of the non-convergence cases and 98.0% of the spurious-attractor cases, clearing 90.9% of TRM's failures.
On multi-solution tasks, 20 samples per problem. 8×8 N-Queens: 99.9% accuracy and 91.4% coverage versus GRAM's 99.7 / 90.3. 10×10: 94.4 vs 89.7 accuracy, 61.5 vs 57.5 coverage. Graph coloring conflicts drop from 2.7 to 0.7 (8 vertices) and from 3.3 to 1.0 (10 vertices), with slightly higher coverage.
The harshest training ablation removes both time conditioning and interpolants: ARC-AGI-1 falls from 58.8 to 43.6, back near TRM. Dropping the decreasing-noise schedule lands at 51.6; dropping noise sharing lands at 56.4. A flow model with no recurrence (FLM) memorizes the training set. Self-conditioning helps generalization; rolling a recurrent state across decreasing noise still wins. SDE versus ODE adds 0.2 to 1.3 points depending on the task.
For anyone training looped reasoners, the practical change is the curriculum: stop asking every step to predict from the same clean input, and walk a noise schedule instead. Same architecture, similar parameter count, fourteen points on ARC-AGI-1. The stability result is more informative than the leaderboard: most of TRM's non-convergence and fake attractors go away.
On problems with many valid answers, different initial noise samples different legal solutions, without GRAM's variational machinery. Training is cheap enough to reproduce: 1–5 hours for Sudoku, 1–2 days for ARC.
This is still a small-model result on structured puzzles, not a recipe for language-model reasoning.
They do not beat FPRM on Maze. 12.2% on ARC-AGI-2 is far from "abstract reasoning works"; ensembling five trajectories only shrinks variance. Data-scarce single-solution tasks need a pseudotarget regularizer: Sudoku trains on 1,000 examples, and later interpolants replace the true solution with the model's previous prediction, otherwise the model overfits.
Training still unrolls k=16 steps. Gradients are truncated, but the forward pass is not free. The authors list simulation-free training as future work. Shared noise admits a linear shortcut that recovers the answer from adjacent interpolants. They argue extra inference steps still help, which rules out "the model only repeats its first guess," but they never probe whether the recurrent state caches the noise. There is no comparison with chain-of-thought language models.