PaperlayerView exact revision on arXiv

Source-linked AI summary

Unrolled Flow Models for Reasoning

Faissal Izermine, Hanru Bai, Oscar Davis, T. Konstantin Rusch

arXiv:2610.09759v1cs.LGcs.AI

TL;DR

Whether additional flow-integration steps improve reasoning remains unclear, because standard objectives supervise sampled time points rather than successive model-generated states. The paper proves a graph-reachability construction and trains through latent rollouts; UFM improves with additional steps, while sphere retraction supports longer reasoning rollouts and multiple rollouts can improve reliability.

  • Problem

    It is unclear whether additional integration steps improve reasoning, and standard flow objectives do not train successive model updates to build on one another.

    Source: §1 · p. 1, §3.2 · p. 5

  • Method

    The paper constructs a two-layer Transformer flow for graph reachability and trains UFM through its own latent rollout with a terminal loss.

    Source: §1 · pp. 1–2

  • Results

    UFM improves with additional steps on reasoning tasks, and an 8.4M-parameter model outperforms 28.6M standard-flow baselines on Sudoku and Maze rollouts.

    Source: §1 · p. 2, §8 · p. 9

  • Takeaways & Limitations

    The results show that rollout training, sphere retraction, and rollout selection help realize reasoning with flows in the evaluated tasks.

    Source: §1 · p. 2, §8 · p. 9

    The experiments use 8M–16M-parameter models on synthetic structured tasks, and selection is less reliable for long Maze-Hard outputs while UFM remains behind the strongest Sudoku-Extreme recursive reasoners.

    Source: §7 · p. 9

Abstract

from arXiv · show

Flow matching enables language generation in few steps, but whether additional integration steps improve reasoning remains unclear. We prove that a flow parameterized by a two-layer Transformer can solve graph reachability, with the required number of integration steps increasing with the target's distance from the root. Yet, standard flow language models can fail to benefit from additional steps on reasoning tasks. We attribute this limitation to objectives that supervise each time point independently, without explicitly training successive steps to build on one another. To address this, we instead train through the model's own latent rollout over a randomly sampled subinterval of [0, 1], decoding only at the endpoint. On ProsQA, this raises accuracy to 97% and enables performance to improve with additional integration steps. For the longer rollouts required by reasoning tasks such as Sudoku and Maze, retracting the latent state onto a sphere stabilizes the dynamics and yields substantial gains over baselines with more than three times as many parameters. Sampling multiple rollouts further improves performance when paired with a parameter-free selection score, although reliable selection remains challenging for longer answers. Together, these results establish a theoretical basis for reasoning with flows and show how rollout training, stable latent dynamics, and rollout selection help realize this capacity in practice.

1 INTRODUCTION

Flow models offer adjustable inference-time computation, but additional steps help reasoning only when updates build on earlier states. This paper gives a graph-reachability construction and proposes rollout training and stabilization methods that make flow steps more useful in experiments.

1 INTRODUCTION

  • A flow’s adjustable number of latent-state updates makes it a candidate for adaptive test-time computation.Unlike autoregressive generation, it repeatedly updates a continuous state and decodes it into a token sequence.Source: §1 · p. 1
  • Additional integration steps do not inherently deepen reasoning; later updates must be able to use information produced by earlier ones.The introduction contrasts this need with flow language models trained for one- or few-step generation.Source: §1 · p. 1
  • A two-layer Transformer flow can solve graph reachability, with Euler steps propagating one hop each and the required count depending on answer distance.This is an existence result and does not establish that training recovers the construction.Source: §1 · pp. 1–2
  • UFM trains through its own latent rollout with a terminal prediction loss, and the paper reports improved reasoning with additional steps on ProsQA and Sudoku-Hard.Sphere retraction stabilizes longer rollouts; the paper also connects the learned dynamics to a superposition-based mechanism.Source: §1 · p. 2
  • Multiple rollouts can further improve reliability, but selecting among them remains challenging for long-answer tasks such as maze solving.The paper’s contribution summary also reports stronger scaling with additional steps than substantially larger flow baselines on structured reasoning tasks.Source: §1 · p. 2

2 BACKGROUND

The background describes conditional flow matching as transporting noise toward an answer through Euler updates, while standard training supervises denoising at a sampled time. It contrasts token-flow baselines with UFM’s latent state and endpoint decoding.

2.1 CONDITIONAL FLOW MATCHING

  • Flow matching defines a path from noise to data and uses a denoiser to predict the clean endpoint from a point on that path.Euler steps approximate integration of the velocity field from t = 0 to t = 1.Source: §2.1 · p. 2
  • A bidirectional DiT receives the flow time as input while prompt tokens stay clean and only answer slots evolve.Standard training supervises a denoising prediction at one sampled time, not a state produced by the model’s preceding update.Source: §2.1 · p. 2

2.2 FLOW LANGUAGE MODELS

  • FLM represents answers as one-hot vectors and moves a vocabulary-sized state toward per-position token distributions at each Euler step.In this paper, FLM denotes the multi-step model rather than the distilled one-step flow-map variant.Source: §2.2 · pp. 2–3
  • S-FLM instead uses learned unit-norm token embeddings and geodesic updates, but still trains a cross-entropy denoiser at one sampled point on its noise path.It also uses a truncated, adaptive noise schedule.Source: §2.2 · p. 3
  • UFM evolves a continuous latent state, trains its rollout jointly with a terminal loss, and applies its vocabulary decoder only at the endpoint.This contrasts with FLM and S-FLM’s sampled-point denoiser training.Source: §2.2 · p. 3

3 INTEGRATION STEPS AND REASONING DEPTH

A two-layer Transformer construction solves graph reachability by propagating information through Euler steps, whose number must cover the target’s distance from the root. In experiments, standard flow baselines gain little from extra steps, while rollout training makes them useful for reasoning.

3.1 A TWO-LAYER FLOW FOR GRAPH REACHABILITY

  • The two-layer construction uses a prompt-caching layer and a linear-attention layer to propagate graph information in a latent answer state.The resulting exact flow separates the reachable candidate from the unreachable one on every valid instance in the theorem’s scope.Source: §3.1 · pp. 3–4
  • Its Euler discretization reaches vertices one hop per step, so correctness requires N at least the distance from the root to the reachable candidate.Using N at least the root’s eccentricity suffices for every valid candidate pair for that root.Source: §3.1 · pp. 4–5
  • The construction is an existence result using orthonormal embeddings and a specialized propagation head; it does not show that gradient descent recovers it.The trained models instead use learned embeddings and reapply their DiT at every step.Source: §3.1 · p. 5

3.2 INTEGRATION STEPS ALONE DO NOT PRODUCE REASONING

  • UFM test accuracy on ProsQA rises from 12% at N = 1 to approximately 97% with longer rollouts, while FLM and S-FLM remain nearly flat.The plotted UFM was trained with Ntrain = 5; the baselines gain little from additional inference steps on the reported reasoning tasks.Source: Figure 2 · p. 5, §3.2 · p. 5
  • Standard flow objectives train denoisers on ground-truth interpolants at sampled times, not on states created by their own preceding updates.The objective therefore does not assess whether one update leaves a state that later updates can use.Source: §3.2 · p. 5
  • UFM samples a rollout grid, evolves latent answer states, and decodes once at the endpoint for its terminal answer loss.The training algorithm stops gradients through early updates when backpropagation is limited to a suffix of the rollout.Source: §3.2 · p. 6

4 UNROLLED FLOW MODEL

UFM trains a continuous answer-state rollout with terminal supervision and decodes only at the endpoint; its intermediate logits show a qualitative reachability pattern.

4 UNROLLED FLOW MODEL

  • UFM trains through its own latent rollout with loss only at the final answer, so successive updates are trained through their combined effect.Training starts from a partially noised answer and integrates over a sampled subinterval; inference starts from noise and integrates the full interval [0, 1].Source: §4 · p. 6
  • A DiT predicts a latent target, and the vocabulary decoder is applied only to the final state rather than intermediate states.Softmax is used on ProsQA and StableMax on the longer Sudoku and Maze tasks.Source: §4 · p. 6

4.1 WHAT THE INTERMEDIATE STATES CARRY

  • Reachable vertices receive higher logits than unreachable ones, and the optimal node separates most within the reachable group as time grows.The pattern qualitatively resembles the construction but does not establish that the trained model implements it.Source: §4.1 · pp. 6–7, Figure 3 · p. 7

5 SCALING TO LONG ROLLOUTS

UFM uses sphere retraction for longer rollouts and outperforms flow baselines on Sudoku-Hard, Sudoku-Extreme, and Maze-Hard; margin-based selection remains less reliable for long outputs.

5 SCALING TO LONG ROLLOUTS

  • The evaluated tasks require longer answer sequences and rollouts than ProsQA.Sudoku-Hard uses 48k training and 2k reporting puzzles; Sudoku-Extreme is evaluated on 422,786 puzzles, and Maze-Hard uses 1,000 training and 1,000 test mazes without a separate validation split.Source: §5 · p. 7, §5.1 · p. 7

5.2 TRAINING LONG ROLLOUTS

  • For long tasks, position-wise sphere retraction limits latent-state growth; training backpropagates through only the final 6 of 24 updates to fit memory.Each training example uses a newly sampled increasing time grid ending at t = 1.Source: §5.2 · pp. 7–8

5.3 COMPARISON WITH FLOW BASELINES

  • 86.9 ± 0.9% on Sudoku-Hard, 74.4 ± 0.0% on Sudoku-Extreme, and 89.3 ± 0.5% on Maze-Hard; UFM exceeds both flow baselines.On Sudoku-Hard, UFM uses 8.4M rather than 28.6M parameters.Source: Table 1 · p. 8, §5.3 · p. 8

5.4 SELECTING AMONG ROLLOUTS

  • K=100 Gaussian-start rollouts are ranked by mean top-two raw-logit margin; on Maze-Hard, 900-cell scoring dilutes signal from roughly 110 path cells and selection trails Pass@100.The margin score is an output-space analogue of δ in Theorem 1.Source: §5.4 · p. 8

6 RELATED WORK

Prior work studies latent reasoning, recurrent networks, and flow-based reasoning; UFM differs in training successive flow updates jointly with an endpoint loss.

6 RELATED WORK

  • Recursive reasoners reuse small weight-tied networks; HRM uses slow and fast modules, TRM a recurrent block and answer state, while EqR and GRAM explore depth and breadth.PTRM also studies sampling and selection.Source: §6 · p. 8
  • Coconut feeds hidden states back as input embeddings, while Zhu et al. (2025) use continuous thoughts to maintain and expand reachable-vertex superpositions.UFM instead updates one latent answer state in place and uses answer-only supervision throughout training.Source: §6 · p. 9
  • DRaFT backpropagates differentiable rewards through diffusion sampling and truncates gradients; UFM instead trains from scratch with endpoint cross-entropy on reasoning tasks.UFM's Nback is analogous to a truncated sampler-backpropagation horizon.Source: §6 · p. 9
  • FRM uses recurrent self-conditioning and reports that more standard FLM sampling steps do not improve Sudoku; Looped Flows use local denoising losses and stop gradients between updates.UFM applies one terminal loss and backpropagates through the final Nback updates.Source: §6 · p. 9

7 LIMITATIONS

The paper’s evidence is limited by model scale, memory constraints, and a gap between its constructive theory and trained dynamics.

7 LIMITATIONS

  • Experiments use 8M–16M-parameter models on synthetic structured tasks, and rollout backpropagation memory limits model size.Larger models and less structured language reasoning require more memory-efficient joint training of consecutive updates.Source: §7 · p. 9
  • The reachability construction assumes orthonormal embeddings and a linear propagation head, and does not show that gradient descent recovers its velocity field.Trained models exhibit only a qualitative pattern in common with the construction.Source: §7 · p. 9
  • Margin-based selection is less reliable for long Maze-Hard outputs, and UFM remains behind the strongest reported recursive reasoners on Sudoku-Extreme.A learned correctness or halting head could improve selection, according to the authors.Source: §7 · p. 9

8 CONCLUSION

The paper presents a theoretical construction for graph reachability and practical methods for making additional flow steps useful in reasoning. Its appendices provide construction details, proofs, rollout procedures, and experimental information.

8 CONCLUSION

  • UFM trains through its own latent rollout, and sphere retraction stabilizes longer rollouts; multiple rollouts improve reliability, though selection is less effective for long answers.The conclusion also reports that an 8.4M-parameter model outperforms 28.6M standard-flow baselines on Sudoku and Maze.Source: §8 · p. 9

APPENDIX OVERVIEW

  • The appendix collects proofs, training and sampling algorithms, and additional experimental details and results.It covers graph reachability, sphere retraction, long rollouts, and experiments.Source: APPENDIX OVERVIEW · p. 13

A.1 NOTATION AND ASSUMPTIONS

  • The construction assumes directed graphs and orthonormal token embeddings, while trained models instead use learned embeddings.The graph may contain cycles, and the construction's embedding assumption requires embedding dimension at least the vocabulary size.Source: §A.1 · p. 13

A GRAPH REACHABILITY: CONSTRUCTION AND PROOFS

  • A two-layer Transformer construction uses Euler flow updates to propagate information along graph edges, with the required steps depending on answer distance from the root.The linear propagation head reads the state as a query and edge buffers as keys and values, then writes its output into the content slot.Source: §8 · p. 9, §A.3 · p. 14, §A.6 · p. 16

A.3 OVERVIEW

  • The first Transformer layer caches prompt information and initializes the answer state; subsequent flow updates use the current state as the propagation query.The prompt remains in the cache, while the velocity is linear in the state.Source: §A.3 · p. 14

A.4 TRANSFORMER PRELIMINARIES

  • Sinusoidal positional encodings let attention select a unique offset position, supporting the construction's attention-chooser heads.The positional encoding separates positions so the score is maximized at the intended offset.Source: §A.4 · pp. 14–15

A.5 LAYER 1: CACHING

  • The first layer's chooser heads cache each edge's endpoints and the two candidate embeddings, and initialize the answer state with the root.At the answer position, the state is initialized as z0 = ϕ⟨A⟩ + ϕr.Source: §A.5 · pp. 15–16

A.6 LAYER 2: LINEAR PROPAGATION

  • The second layer uses unnormalized linear attention to read edge buffers and define the velocity field, without a softmax or normalizing denominator.The head reads the content slot as query, buffer 1 as key, and buffer 2 as value, then writes back to the content slot.Source: §A.6 · p. 16

A.7 PROOF OF THEOREM 1

The proof connects the constructed flow's vertex scores to directed walks from the root, then shows when Euler discretization separates a reachable candidate. It also describes noisy initialization, full-vocabulary readout, and the local effect and boundary of sphere retraction.

A.7 PROOF OF THEOREM 1

  • For t > 0, the exact flow assigns a positive vertex coordinate exactly to vertices reachable from the root, with a positive lower bound for reachable vertices.The proof uses the matrix exponential of the graph adjacency operator; the shortest path supplies a positive walk-count term.Source: §A.7 · p. 17

A.8 EULER DISCRETIZATION

  • Euler updates separate the reachable candidate once N ≥ dist(r, ci⋆); for smaller N, both candidate scores are zero and the readout does not separate them.The discretized coefficient is positive exactly when a walk of length at most N connects the root to a vertex.Source: §A.8 · p. 17

A.9 NOISY INITIALIZATION

  • With noisy initialization, noise orthogonal to the vertex span is frozen and invisible to readout, while the expected state follows the noiseless flow.This claim is stated for the noisy initialization considered in the proof.Source: §A.9 · p. 18

A.10 READOUT OVER THE FULL VOCABULARY

  • A fixed linear decoder using the content slot and candidate buffers returns the reachable candidate with margin at least δ over the full vocabulary.The decoder adds a candidate-specific score offset while leaving the velocity field unchanged.Source: §A.10 · p. 18

B.1 LOCAL GEOMETRIC EFFECT

  • Sphere retraction makes the leading-order update tangent to the sphere, preserving norm in continuous time; the analysis does not establish convergence through t = 1.The implemented-update interpretation requires local Lipschitz continuity and an interval bounded away from the endpoint.Source: §B.1 · pp. 18–19

B.2 RETRACTION ABLATION AND NORM DIAGNOSTIC

The retraction ablation favors retraction on Sudoku-Extreme, while the norm diagnostic shows that it limits growth without constraining the raw state to a fixed norm.

B.2 RETRACTION ABLATION AND NORM DIAGNOSTIC

  • 67% test accuracy with retraction versus 55% without retraction at Neval = 64 in separately trained Sudoku-Extreme models.The ablation used a different recipe from the main cosine-schedule comparison.Source: §B.2 · p. 19
  • Without retraction, the raw latent norm grows roughly 44,000 times its initial scale over a 252-step rollout; retraction limits growth but does not fix the raw-state norm.The fixed radius applies to the retracted update base, not necessarily the raw state.Source: §B.2 · p. 19
  • The reachability construction also permits latent norm growth, while positive rescaling leaves its final candidate decision unchanged.This is a property of the construction's final decision, not the Sudoku-Extreme ablation.Source: §B.2 · p. 19

C LONG-ROLLOUT TRAINING AND INFERENCE

Long-rollout UFM training samples a time grid, updates the latent answer state through a rollout, and applies a prediction loss at the endpoint. The appendix also specifies inference and the reporting convention for single-rollout accuracy.

C LONG-ROLLOUT TRAINING AND INFERENCE

  • Training samples a time grid ending at 1, rolls the latent answer state forward, decodes at the endpoint, and backpropagates through only the final Nback steps.The initial state interpolates noise and the embedded answer at the sampled starting time.Source: C LONG-ROLLOUT TRAINING AND INFERENCE · p. 20
  • During long-rollout inference, the denoiser reads the raw state, while retraction is applied to the update base after the first step.The decoded output is the highest-probability token at each position.Source: C LONG-ROLLOUT TRAINING AND INFERENCE · p. 20

D EXPERIMENTAL DETAILS

  • Single-rollout accuracy is reported as exact-match mean and standard deviation over five noise seeds, using a fixed checkpoint and a uniform Neval = 128 grid for Table 3.Reported values are rounded to one decimal place, and the checkpoints disable the DiT time input.Source: D EXPERIMENTAL DETAILS · p. 21, Table 3 · p. 21

D.1 SELECTING AMONG ROLLOUTS

With 100 rollouts, margin selection nearly matches Pass@100 on the Sudoku-Extreme subset but trails it on Maze-Hard; the appendix also details benchmark protocols, baseline comparisons, and implementation settings.

D.1 SELECTING AMONG ROLLOUTS

  • 98.6% margin-based selection on the Sudoku-Extreme subset nearly matches 98.65% Pass@100, while Maze-Hard selection reaches 92.0% versus 95.0% Pass@100.Each problem uses K = 100 rollouts; the Sudoku evaluation uses a 2000-puzzle subset and Maze-Hard uses all 1000 test mazes.Source: §D.1 · p. 21
  • The parameter-free chooser ranks rollouts by the mean per-position gap between the two largest raw logits, an output-space analogue of Theorem 1's separation margin.The logits are measured before probability normalization.Source: §D.1 · p. 21
  • Margin selection is less reliable on long Maze-Hard answers: averaging over 900 cells can dilute path differences, unlike 81-position Sudoku answers.On 18 recoverable boards, correct and incorrect rollouts differed in mean margin by 0.07 around mean margins of 200, and the chooser selected a correct rollout on 9 of 18 boards versus 38% for a uniform choice among distinct answers.Source: §D.1 · p. 21

D.2 PROSQA

  • ProsQA evaluation uses questions whose answers are at most four hops from the root; UFM trains with five steps while baselines see about 22.3M examples versus UFM's 4.4M.The table captions distinguish single-trajectory accuracy from multiple-rollout results, and published recursive-reasoner entries are not reruns.Source: §D.2 · pp. 21–22, Table 4 · p. 22, Table 5 · p. 22

D.3 COMPARISON WITH RECURSIVE REASONERS

  • Published recursive-reasoner results provide context rather than a controlled pairwise comparison with UFM.Source: §D.3 · p. 22

D.4 BASELINE REPLICATION ON SUDOKU-HARD

  • Baseline checkpoints were selected on the same 2k Sudoku-Hard puzzles reported, favoring baselines, while UFM was selected on a separate validation split; the updated UFM result is 86.9 ± 0.9% at N = 128.FLM and S-FLM used their best 30k-step checkpoints; the updated UFM result averages five noise seeds.Source: §D.4 · p. 23, Table 6 · p. 23

D.5 HYPERPARAMETERS

  • The Sudoku and Maze UFM backbone is a two-layer DiT with width 448, eight heads, and 8.4M parameters; reported checkpoints disable sinusoidal time conditioning.Training uses 24 updates, backpropagates through the final six, samples the initial time and subsequent times, and applies one endpoint StableMax cross-entropy loss.Source: §D.5 · p. 23, Table 7 · p. 24

D.6 COMPUTING RESOURCES

  • UFM training used one A100 for Sudoku-Hard and Sudoku-Extreme and two for Maze-Hard; Sudoku-Hard training took 1.14 hours for 19.7k optimization steps.The baseline runs took 1.07 and 1.12 hours for 20k steps, while selected 30k-step checkpoints took 1.61 and 1.68 hours.Source: §D.6 · p. 24

Related papers

Loading 2610.09759v1…