Paper deep dive
Unrolled Flow Models for Reasoning
Faissal Izermine, Hanru Bai, Oscar Davis, T. Konstantin Rusch
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 10/11/2026, 3:07:10 AM
Summary
This paper introduces Unrolled Flow Models (UFM) to enhance reasoning capabilities in flow-based language models. Standard flow models often fail to leverage additional integration steps for complex reasoning tasks because they supervise each time point independently. UFM addresses this by training through the model's own latent rollout with a terminal loss, decoding only at the endpoint. Theoretical proofs show a two-layer Transformer can solve graph reachability via flow integration. Empirical results on ProsQA, Sudoku, and Maze tasks demonstrate that UFM, combined with sphere retraction for stability and multiple rollout sampling, significantly outperforms baselines like FLM and S-FLM, achieving high accuracy and scaling with integration steps.
Entities (10)
Relation Signals (8)
Unrolled Flow Model → improvesaccuracyon → ProsQA
confidence 95% · On ProsQA, this raises accuracy to 97% and enables performance to improve with additional integration steps.
Two-layer Transformer → solvestask → Graph Reachability
confidence 95% · We prove that a flow parameterized by a two-layer Transformer can solve graph reachability
Unrolled Flow Model → outperforms → FLM
confidence 92% · UFM improves with additional steps, whereas evaluated flow baselines show little gain.
Unrolled Flow Model → outperforms → S-FLM
confidence 92% · UFM improves with additional steps, whereas evaluated flow baselines show little gain.
Unrolled Flow Model → appliedto → Sudoku
confidence 90% · For the longer rollouts required by reasoning tasks such as Sudoku and Maze
Unrolled Flow Model → appliedto → Maze
confidence 90% · For the longer rollouts required by reasoning tasks such as Sudoku and Maze
Unrolled Flow Model → usestechnique → Sphere Retraction
confidence 90% · For the longer rollouts required by reasoning tasks such as Sudoku and Maze, retracting the latent state onto a sphere stabilizes the dynamics
Sphere Retraction → stabilizes → Latent Dynamics
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2610.09759v1
- Canonical: https://arxiv.org/abs/2610.09759v1
Trouble viewing inline? Open PDF directly →
Full Text
88,238 characters extracted from source content.
Expand or collapse full text
Unrolled Flow Models for Reasoning Faissal Izermine Hanru Bai Affiliation: Max Planck Institute for Intelligent Systems, ELLIS Institute Tübingen, ETH Zurich, Oscar Davis T. Konstantin Rusch Affiliation: University of Oxford, Tübingen AI Center, Liquid AI Abstract 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][0,1], decoding only at the endpoint. On ProsQA, this raises accuracy to 97%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. The implementation can be found here. 1 Introduction Autoregressive language models can spend more computation at inference by generating a longer chain of thought (Wei et al., 2022) or, more broadly, by scaling test-time compute (Snell et al., 2024). Flow matching (Lipman et al., 2023; Liu et al., 2023; Albergo et al., 2025) offers a different way to produce discrete answers (Eijkelboom et al., 2024). Instead of predicting one token after another, flow models start from noise and repeatedly update a continuous state that is decoded into a token sequence. The number of updates can be chosen at inference time. This makes flow models an appealing candidate for adaptive test-time computation. Additional integration steps do not inherently guarantee deeper reasoning. In generation tasks, a small number of steps is often enough, and recent work has pushed flow language models toward one- or few-step sampling by learning flow maps that make direct jumps along a trajectory (Lee et al., 2026; Roos et al., 2026; Davis et al., 2026; Yoo et al., 2026). For reasoning, additional steps are useful only if later updates can build on information produced by earlier ones. Prior work in latent and recursive reasoning motivates this work. Coconut shows that autoregressive language models can update a continuous latent state instead of emitting every intermediate thought as text (Hao et al., 2026). On graph reachability, Zhu et al. (2025) show that repeated latent updates can maintain a superposition of reachable vertices and expand it one hop at a time. Recursive reasoners similarly solve structured tasks by applying a small, weight-tied network repeatedly to an internal state (Wang et al., 2025; Jolicoeur-Martineau, 2025). An Euler rollout of a flow has the same basic form: one network is applied again and again to a latent state. We investigate how these repeated applications cooperate. To show what integration steps can compute, we construct explicit weights for a two-layer Transformer that defines a flow solving the graph reachability problem. Its Euler rollout reaches vertices one hop per step, so the number of steps needed to separate the candidates depends on the distance from the root to the answer. However, this existence result does not establish that training recovers the construction (Section 3.1). Standard training of flow models does not directly assess whether an update produces a state that later updates can leverage. To address this limitation, we propose Unrolled Flow Model (UFM), which trains through its own latent rollout using a single prediction loss at the final state. To reduce memory usage for long rollouts, we backpropagate this loss through only the last few steps (Section 4). We stabilize these rollouts using sphere retraction to control latent-state growth (Section 5). On ProsQA and Sudoku-Hard, UFM improves with additional steps, whereas evaluated flow baselines show little gain. Multiple rollouts further improve reliability, although selecting among them remains challenging for tasks with long answers, such as maze solving. Our contributions are: (i) We give a constructive result showing how flow integration can implement graph search with the number of steps determined by the instance. (i) We introduce training through a terminal rollout loss and show that it makes flow steps useful for reasoning in our experiments, while connecting the resulting dynamics to a superposition-based mechanism. (i) We use sphere retraction to stabilize long rollouts and demonstrate stronger scaling with additional steps than substantially larger flow baselines on structured reasoning tasks. 2 Background 2.1 Conditional flow matching Setup. Given a prompt c, the model generates an answer with L positions. A flow transforms noise at t=0t=0 into an answer representation at t=1t=1, using N Euler steps on a grid 0=t0<⋯<tN=10=t_0<·s<t_N=1. Flow matching. Flow matching (Lipman et al., 2023; Liu et al., 2023; Albergo et al., 2025) defines a path between noise 0∼p0=(0,) x_0 p_0=N(0, I) and data 1∼p1 x_1 p_1. We use the linear interpolant It:=(1−t)0+t1,^θ(,t)≈[1∣It=].I_t:=(1-t)\, x_0+t\, x_1, x_θ( x,t) [ x_1 I_t= x ]. (1) The denoiser ^θ x_θ predicts the clean endpoint from a point on this path. Since 1−0=(1−It)/(1−t) x_1- x_0=( x_1-I_t)/(1-t), its corresponding velocity field is θ(,t)=^θ(,t)−1−t. v_θ( x,t)= x_θ( x,t)- x1-t. (2) Integrating ˙t=θ(t,t) z_t= v_θ( z_t,t) from t=0t=0 to t=1t=1 transports noise toward the data distribution. In practice, we approximate this integration with Euler steps: tk+1=tk+(tk+1−tk)θ(tk,tk). z_t_k+1= z_t_k+(t_k+1-t_k)\, v_θ( z_t_k,t_k). (3) Conditioning and training. The prompt tokens are concatenated with L answer slots and processed by a bidirectional Diffusion Transformer (DiT) (Vaswani et al., 2017; Peebles and Xie, 2023). The DiT receives the flow time t as a conditioning input, while the prompt remains clean throughout the rollout; only the answer slots evolve. The conditional denoiser is trained using: ℒFM(θ)=(c,1),0,t∼[0,1][ℓ(^θ(It,t,c),1)].L_FM(θ)=E_(c, x_1),\, x_0,\,t [0,1] [ ( x_θ(I_t,t,c),\ x_1 ) ]. (4) Here, ℓ is the per-example loss. Thus, each training example supervises a denoising prediction at one sampled time; the model is never evaluated on a state produced by its own preceding update. We return to this distinction in Section 3.2. 2.2 Flow Language Models Flow Language Models (FLMs) apply this formulation to discrete tokens. FLM (Lee et al., 2026) represents the clean answer as a stack of one-hot vectors, 1∈0,1L×V x_1∈\0,1\^L× V. Its denoiser therefore predicts a distribution over the vocabulary at every answer position: ^θ(,t,c)=(pi(⋅∣It=,c))1≤i≤L. x_θ( x,t,c)= (p^i(· I_t= x,\ c) )_1≤ i≤ L. (5) The model is trained with cross-entropy, and the evolving state t∈ℝL×V z_t ^L× V is moved toward this predicted distribution at each Euler step. In our experiments, FLM refers to the multi-step flow model, rather than the distilled one-step flow-map variant introduced in the same work. S-FLM. S-FLM (Deschenaux and Gulcehre, 2026) instead represents tokens with learned unit-norm embeddings on the hypersphere. It follows a geodesic path from noise toward these embeddings, rather than evolving a vocabulary-sized state as in FLM. Its denoiser still produces a distribution over the vocabulary; each sampling step follows a posterior-weighted average of geodesic directions toward the token embeddings. Both methods train a cross-entropy denoiser at one sampled point on their respective noise paths. S-FLM additionally uses a truncated, adaptive noise schedule. ELF (Hu et al., 2026) also defines the flow in a continuous token-embedding space. FLM and S-FLM are our flow baselines. UFM also evolves a continuous latent state, but it predicts a latent target and applies the vocabulary decoder only once, at the endpoint. Its rollout is trained jointly with a terminal loss, as described in Section 4. 3 Integration Steps and Reasoning Depth 3.1 A Two-Layer Flow for Graph Reachability We define the graph-reachability task and the flow construction used in the existence result. The theorem below shows that the construction solves the task; its Euler discretization will later show how many integration steps are needed. Notation is collected in Table 2. Graph and prompt. Let =(,ℰ)G=(V,E) be a directed graph with n=||≤nmaxn=|V|≤ n_ vertices and m=|ℰ|m=|E| edges. Let r∈r be a designated starting vertex, called the root, and let c1,c2c_1,c_2 be two candidate destinations. Exactly one candidate, denoted ci⋆c_i , is reachable from r; we write c3−i⋆c_3-i for the other candidate. We write dist(r,v)dist(r,v) for the length of a shortest directed path from r to v, and k:=v∈:dist(r,v)≤kV_k:=\v :dist(r,v)≤ k\ for the vertices reachable within k hops. The sequence 0⊆1⊆⋯V_0 _1 ·s stops growing at the eccentricity of the root, D:=maxdist(r,v):v is reachable from r.D:= \dist(r,v):v is reachable from r\. Thus, dist(r,ci⋆)≤Ddist(r,c_i )≤ D. We also write Wk(r,v)W_k(r,v) for the number of directed walks of length k from r to v. Following Zhu et al. (2025), the prompt lists the directed edges and then asks which candidate is reachable (Figure 4 in the appendix): c=(⟨s⟩,u1,v1,⟨e⟩,…,um,vm,⟨e⟩,⟨Q⟩,c1,c2,⟨R⟩,r,⟨A⟩).c=( ,u_1,v_1, ,…,u_m,v_m, , ,c_1,c_2, ,r, ). The flow state t z_t evolves at the position marked by ⟨A⟩ , where the final answer is decoded. Representation. Let V be the vocabulary size and dted_te the embedding dimension. Each token v has an embedding ϕv∈ℝdte φ_v ^d_te. As in Zhu et al. (2025), we assume orthonormal token embeddings, ⊤=,=[ϕ1,…,ϕV],dte≥V. = I, =[ φ_1,…, φ_V], d_te≥ V. This lets ⟨ϕv,t⟩ φ_v, z_t recover the coefficient of vertex v in the state. This assumption applies solely to the construction; the trained models use standard embeddings. The first layer stores each edge’s source and target embeddings in separate buffers at its ⟨e⟩ position, and the two candidate embeddings in buffers at ⟨A⟩ . It initializes the answer state as 0=ϕ⟨A⟩+ϕr z_0= φ_ + φ_r (Lemmas 6 and 8). Constructed velocity field. The second layer is a softmax-free linear-attention head (Definition 9). It reads the content slot of ⟨A⟩ as its query and the cached source and target edge embeddings as its keys and values. Using these cached edge endpoints, it applies =∑(uj→vj)∈ℰϕvjϕuj⊤ A_G= _(u_j→ v_j) φ_v_j φ_u_j , the adjacency operator of G in embedding coordinates. With a fixed gain λ>0λ>0, the head induces ˙t=λt,0=ϕ⟨A⟩+ϕr. z_t=λ A_G z_t, z_0= φ_ + φ_r. (6) We use a linear head where Zhu et al. (2025) use softmax attention because it makes the velocity linear in the state, so the flow has a closed form. At time t=1t=1, the construction compares ⟨ϕc1,1⟩ φ_c_1, z_1 and ⟨ϕc2,1⟩ φ_c_2, z_1 and returns the candidate with the larger score. This is a candidate-restricted readout. The experiments use the stricter full-vocabulary decoder; we discuss that difference in Remark 11. Exact flow The constructed velocity field has a closed-form solution. This gives the existence result: the flow separates the reachable candidate from the unreachable one on every instance. GaussianNoise(0,d)N(0, I_d)Prompt ccz0(1)z_0^(1)z0(2)z_0^(2)⋮ 0(L)z_0^(L)Time 00z0∈ℝL×dz_0 ^L× dPrompt cczt(1)z_t^(1)zt(2)z_t^(2)⋮ (L)z_t^(L)Time ttzt∈ℝL×dz_t ^L× dPrompt ccz1(1)z_1^(1)z1(2)z_1^(2)⋮ 1(L)z_1^(L)Time 11z1∈ℝL×dz_1 ^L× dDecoded→Vd→ V Softmax VocabSpaceℝVR^VLatent FlowLatent Flow Figure 1: UFM computes in latent space and decodes once. At inference, a DiT repeatedly updates a noisy latent answer state from t=0t=0 to t=1t=1. A linear vocabulary decoder is applied only to the final state. Theorem 1 (Existence). Fix a vocabulary of size V, nmaxn_ , and λ>0λ>0. There exists a two-layer Transformer, consisting of a prompt-caching layer followed by a linear-attention layer, whose weights are independent of the graph, root, and candidate pair. Its answer state satisfies Equation 6 for every directed graph with ||≤nmax|V|≤ n_ , every root r, and every candidate pair with exactly one reachable candidate ci⋆c_i . At time t=1t=1, ⟨ϕv,1⟩=∑k≥0λk!Wk(r,v),v∈. φ_v, z_1 = _k≥ 0 λ^kk!\,W_k(r,v), v . This coefficient is positive exactly when v is reachable from r. The candidate-restricted readout therefore returns ci⋆c_i , with the margin bound δ:=min0≤k≤nmax−1λk!≤⟨ϕci⋆,1⟩−⟨ϕc3−i⋆,1⟩.δ:= _0≤ k≤ n_ -1 λ^kk!≤ φ_c_i , z_1 - φ_c_3-i , z_1 . Why the exact flow solves reachability. Since ϕ⟨A⟩= A_G φ_ = 0, the solution is t=ϕ⟨A⟩+eλtϕr z_t= φ_ +e^λ t A_G φ_r. Orthonormality makes kϕr A_G^k φ_r count length-k walks (Lemma 10), giving the theorem’s nonnegative coefficient series. An unreachable vertex has no such walks; a reachable vertex has a shortest path of length k⋆≤nmax−1k ≤ n_ -1, contributing at least λk⋆/k⋆!≥δλ^k /k !≥δ at t=1t=1. The network realization is established in Lemmas 8 and 9. Unlike the Euler rollout considered next, the exact flow assigns positive mass to every reachable vertex at any t>0t>0. Reaching vertices one hop per step is therefore a property of the discretization, not of the flow itself. Corollary 2 (Euler steps match graph distance). For any Euler grid 0=t0<⋯<tN=10=t_0<·s<t_N=1 with positive step sizes, tN=ϕ⟨A⟩+∏i=0N−1(+λΔti)ϕr. z_t_N= φ_ + _i=0^N-1 ( I+λ t_i A_G ) φ_r. Hence, for every vertex v, ⟨ϕv,tN⟩>0⟺v∈N. φ_v, z_t_N >0 v _N. Thus, the candidate-restricted readout is correct when N≥dist(r,ci⋆)N (r,c_i ). Taking N≥D=ecc(r)N≥ D=ecc(r) suffices for every valid candidate pair with this root. Expanding the product in Corollary 2 gives positive coefficients for 0,…,N A_G^0,…, A_G^N. The state therefore contains contributions from walks of length at most N, and its vertex support is exactly NV_N. Step sizes change the weights on these walks and the number of steps determines how far information can propagate: a target farther from the root requires more Euler steps, using the same weights. Remark 3 (Noisy initialization). For 0=ϕ⟨A⟩+ϕr+σ z_0= φ_ + φ_r+σ ξ, with ∼(0,) ξ (0, I), the exact time-11 flow returns the correct candidate with probability at least 1−exp(−δ24σ2e2λm).1- \! (- δ^24σ^2e^2λ m ). The bound weakens with noise and graph size; see Section A.9. The construction uses orthonormal embeddings, a softmax-free propagation head, a prompt cache, and a time-independent velocity field. It does not show that gradient descent recovers this solution. Trained UFM uses learned embeddings and reapplies its DiT at every step; Section 4.1 examines whether its dynamics show a related mechanism. 3.2 Integration Steps Alone Do Not Produce Reasoning We test whether increasing the number of integration steps at inference improves reasoning accuracy in standard flow models. Setup. Figure 2: Rollout training makes additional steps useful. ProsQA test accuracy versus the number of Euler steps at inference. FLM and S-FLM remain nearly flat, while a single UFM trained with Ntrain=5N_train=5 improves from 12%12\% at N=1N=1 to approximately 97%97\% with longer rollouts. Decoding uses the full vocabulary. As in Zhu et al. (2025), we use ProsQA (Hao et al., 2026), a directed graph-reachability task with two candidate answers. The models have two layers and approximately 15.715.7M parameters. The plotted UFM is trained with Ntrain=5N_train=5 steps and gradients through all five updates. UFM and S-FLM have approximately matched training budgets measured by the total number of model evaluations across examples; configurations are given in Appendix D. Figure 2 shows test accuracy as the number of Euler steps N increases at inference. Both baseline curves are nearly flat: more integration steps add little accuracy. We revisit the same baselines on the harder Sudoku and Maze benchmarks in Section 5.3. On Sudoku-Hard, the 28.628.6M baselines change by only about one point from N=8N=8 to N=128N=128 (Table 6). Consequently, standard flow models require more than additional forward passes to improve reasoning. This behavior is consistent with the standard flow objective in Equation 4. FLM and S-FLM train the denoiser to recover the answer from a ground-truth interpolant at a sampled time. The model is therefore not trained on the states created by its own preceding updates, and the objective does not assess whether one update leaves a useful state for the next. UFM instead trains a sequence of updates against one terminal answer loss. S-FLM already evolves an embedding-space state; yet, its accuracy changes little with additional steps. Our construction motivates keeping the computation in latent space and decoding only at the endpoint; Section 4 trains these updates jointly through the rollout. 4 Unrolled Flow Model We want early updates to produce states that later updates can use to reach the answer. We therefore train an Unrolled Flow Model (UFM) through its own rollout with a loss on the final answer. UFM keeps the answer state in a continuous latent space and decodes only once, at the endpoint. Latent rollout. t∈ℝL×d z_t ^L× d, where d is the hidden width. A DiT predicts a latent target, ^θ(t,t,c)∈ℝL×d x_θ( z_t,t,c) ^L× d, and defines the velocity as in Equation 2. Unlike FLM and S-FLM, intermediate states are not decoded into distributions over the vocabulary. The vocabulary decoder is applied only to the final state, ^=softmax(tNdec⊤),dec∈ℝV×d. y=softmax\! ( z_t_N W_dec ), W_dec ^V× d. We use softmax on ProsQA and replace it with StableMax (Prieto et al., 2025) on the longer Sudoku and Maze tasks. Details are in Appendix C. Terminal rollout loss. We apply the loss only at the endpoint, so backpropagation trains successive updates through their combined effect on the final answer. During training, the rollout starts from a partially noised answer at a sampled time t0t_0, t0=(1−t0)+t0Embed() z_t_0=(1-t_0) ξ+t_0Embed( y) with ∼(0,) ξ (0, I), and is integrated over a sampled subinterval ending at tN=1t_N=1. ℒUFM(θ)=(c,),,t0[CE(softmax(tNdec⊤),)].L_UFM(θ)=E_(c, y),\, ξ,\,t_0 [CE\! (softmax( z_t_N W_dec ),\, y ) ]. (7) No intermediate state receives a vocabulary loss. At inference, we set t0=0t_0=0, start from noise, and integrate the full interval [0,1][0,1]. The complete long-rollout training and inference procedures, including sphere retraction and truncated backpropagation, are given in Appendix C. Algorithm 1 Base UFM training on one batch. Task-specific variants are described in the text. 0: Batch (c,)(c, y), parameters θ, rollout length NtrainN_train, backpropagation length 1≤Nback≤Ntrain1≤ N_back≤ N_train 1: Sample noise ∼(0,) ξ (0, I) in the latent answer space 2: Sample a task-specific grid 0≤t0<t1<⋯<tNtrain=10≤ t_0<t_1<·s<t_N_train=1 3: t0←(1−t0)+t0Embed() z_t_0←(1-t_0) ξ+t_0Embed( y) 4: for k=0k=0 to Ntrain−1N_train-1 do 5: Δtk←tk+1−tk t_k← t_k+1-t_k 6: ^θ(k)←^θ(tk,tk,c) x_θ^(k)← x_θ( z_t_k,t_k,c) 7: θ(k)←(^θ(k)−tk)/(1−tk) v_θ^(k)←( x_θ^(k)- z_t_k)/(1-t_k) 8: tk+1←tk+Δtkθ(k) z_t_k+1← z_t_k+ t_k\, v_θ^(k) 9: if k<Ntrain−Nbackk<N_train-N_back then 10: tk+1←sg(tk+1) z_t_k+1 ( z_t_k+1) stop gradients through early updates 11: end if 12: end for 13: ^←softmax(tNtraindec⊤) y ( z_t_N_train W_dec ) decode once 14: ℒ←CE(^,)L ( y, y) 15: return ℒ,∇θℒL, _θL 4.1 What the Intermediate States Carry Figure 3: Intermediate states separate reachable vertices. We plot the logit ⟨t,v⟩ z_t, w_v of vertex v, where v w_v is row v of the final decoder dec W_dec. Logits are centered separately for each problem by their mean over vertices. UFM trained with Ntrain=5N_train=5 is evaluated at t=0.2,0.4,0.6,0.8t=0.2,0.4,0.6,0.8, on the 217217 ProsQA test problems with dist(r,ci⋆)=4dist(r,c_i )=4, averaged over four seeds. Vertices are grouped as not reachable, reachable, frontier (dist(r,v)=kdist(r,v)=k), and optimal (the vertex at hop k on a shortest root-to-answer path). Following Section 3.1, we inspect the logit ⟨t,v⟩ z_t, w_v of each vertex v, where v w_v is its output embedding. In Figure 3, non-reachable vertices remain near zero at every step, while several reachable vertices receive positive logits. At the first step, no single reachable vertex is clearly selected. The reachable group sits above the not-reachable one and drifts further right as t grows, and within it the optimal node separates the most. While it does not establish that the trained model implements the construction, this qualitatively resembles Corollary 2, whose Euler states have support on all vertices reachable within k hops after k steps. 5 Scaling to Long Rollouts We evaluate UFM on Sudoku-Hard, Sudoku-Extreme, and Maze-Hard. These tasks require longer answer sequences and longer rollouts than ProsQA. 5.1 Benchmarks and evaluation protocol Sudoku-Hard. We use the split of Deschenaux and Gulcehre (2026): 4848k training puzzles and 22k reporting puzzles, each with 3030 givens. UFM checkpoints are selected on an additional 22k validation split, disjoint from both sets. Our best FLM and S-FLM checkpoints are selected on the reporting set. Sudoku-Extreme. Following Wang et al. (2025); Jolicoeur-Martineau (2025), training uses 11k base puzzles with 10001000 augmentations each. We evaluate on the full test set of 422,786422,786 puzzles. Checkpoint selection uses 22k additional puzzles from the original training pool, disjoint from the selected 11k base training puzzles and the test set. Each answer contains 8181 positions. Maze-Hard. We use the 30×3030× 30 benchmark of Wang et al. (2025); Jolicoeur-Martineau (2025), with 10001000 training and 10001000 test mazes whose shortest paths exceed 110110 cells. The model predicts a 900900-cell grid, of which roughly 110110 cells lie on the solution path. UFM and both flow baselines select checkpoints on the test set; no separate validation split is used. 5.2 Training long rollouts Sphere retraction. On long Sudoku rollouts, the latent norm can grow sharply (Appendix B). This changes the scale of the inputs to later updates and can destabilize the rollout. We control this growth on Sudoku and Maze with a position-wise sphere retraction, ψ()=d∥2.ψ( z)= d\, z z _2. Let tk z_t_k denote the state read by the DiT and ytk=ψ(tk)y_t_k=ψ( z_t_k) its retracted version. The update is tk+1=ytk+Δtk1−tk(^θ(tk,tk,c)−ytk),ytk+1=ψ(tk+1), for 1≤k<N. z_t_k+1=y_t_k+ t_k1-t_k ( x_θ( z_t_k,t_k,c)-y_t_k ), y_t_k+1=ψ( z_t_k+1), for 1≤ k<N. This is a projected Euler-type update: after the first step, its base is retracted, while the DiT continues to read the raw state. We use this retraction on Sudoku and Maze only. In a separate Sudoku ablation, training without it lowers accuracy; details are in Appendix B. Truncated Backpropagation For the long tasks, we reduce activation memory by backpropagating only through the final Nback=6N_back=6 of Ntrain=24N_train=24 updates, as in Jolicoeur-Martineau (2025); Clark et al. (2024). The first 1818 updates are computed without recording gradients and produce the state refined by the final six. We choose Nback=6N_back=6 to fit the available memory. Each training example uses a newly sampled increasing time grid ending at t=1t=1, exposing the model to varying step sizes. The complete long-rollout procedures are given in Appendix C. 5.3 Comparison with Flow Baselines We compare all methods at Neval=128N_eval=128 on Sudoku-Hard (Deschenaux and Gulcehre, 2026). FLM and S-FLM use the authors’ released code and the best tested checkpoints and final-step settings at this budget (Section D.4). Table 1: Accuracy (%) on Sudoku-Hard. FLM and S-FLM are our runs of the authors’ released code. Train is training-only wall-clock time on one NVIDIA A100. All methods use the same training data. Method Params Train (h) Sudoku-Hard FLM (Lee et al., 2026)a 28.628.6M 1.61.6 51.9 S-FLM (Deschenaux and Gulcehre, 2026)a 28.628.6M 1.71.7 50.9b UFM (ours) 8.48.4M 1.11.1 86.9±0.986.9± 0.9 aSelected on the 22k reporting puzzles; UFM is selected on a separate 22k validation split. The UFM result is the mean and standard deviation over five independent noise seeds at Neval=128N_eval=128 on a uniform grid. bWith ancestral noise removal at Neval=128N_eval=128; with the default greedy final step, S-FLM peaks at 50.2%50.2\% at N=64N=64 (Table 6). On Sudoku-Hard, UFM reaches 86.9±0.9%86.9± 0.9\%, compared with 51.9%51.9\% for FLM and 50.9%50.9\% for S-FLM, using 8.48.4M rather than 28.628.6M parameters. The gap widens on the harder tasks. On Sudoku-Extreme, UFM reaches 74.4±0.0%74.4± 0.0\% with a uniform Neval=128N_eval=128 grid, averaged over five noise seeds. FLM and S-FLM reach 10.7%10.7\% and 9.4%9.4\%, respectively. On Maze-Hard, UFM reaches 89.3±0.5%89.3± 0.5\%, compared with 40.4%40.4\% for FLM and 49.9%49.9\% for S-FLM (Table 4). Published recursive-reasoner results and FRM are included for context in Tables 4 and 5. 5.4 Selecting Among Rollouts Inference begins from Gaussian noise, so different rollouts can produce different answers. We therefore draw K=100K=100 rollouts and select one with a parameter-free score. Margin averages the gap between the two largest raw logits over output positions. It is an output-space analogue of δ in Theorem 1. On mazes, scores average over 900900 cells, although only roughly 110110 describe the solution path. This makes selection harder, and selected accuracy remains below Pass@100. Detailed diagnostics are in Section D.1. 6 Related work Recursive and looped reasoning. Recursive reasoners repeatedly apply a small, weight-tied network to a latent state. HRM (Wang et al., 2025) uses slow and fast recurrent modules, while TRM (Jolicoeur-Martineau, 2025) uses a single recurrent block and an answer state. Subsequent work improves stability and scaling (Movahedi et al., 2026; Fein-Ashley and Rashidinejad, 2026; Gao et al., 2025). EqR and GRAM explore both depth, through more iterations, and breadth, through multiple trajectories (Huang et al., 2026; Baek et al., 2026). PTRM also studies sampling and selection (Sghaier et al., 2026). Latent chain of thought. Coconut (Hao et al., 2026) feeds a language model’s hidden state back as its next input embedding and introduces ProsQA. Zhu et al. (2025) show that a two-layer Transformer with continuous thoughts can solve graph reachability by maintaining a superposition of reachable vertices and expanding it over several steps. Their result provides the starting point for our construction in Section 3.1. UFM updates one latent answer state in place and uses answer-only supervision throughout training. Coconut appends continuous thoughts and uses a curriculum that gradually replaces language reasoning steps. Differentiating through generation. Backpropagating a terminal objective through a generative sampler is established in diffusion reward fine-tuning. DRaFT (Clark et al., 2024) backpropagates differentiable rewards through a diffusion sampler and introduces DRaFT-K, which truncates gradients to the final K sampling steps. UFM uses the same general optimization mechanism, but is trained from scratch with endpoint cross-entropy on reasoning tasks rather than fine-tuned against an external reward. Our NbackN_back plays the role of a truncated sampler-backpropagation horizon. Flow models for reasoning. Flow Reasoning Models (FRM) (Helbling et al., 2026) add recurrent self-conditioning to a flow model. Their Fixed-Point Forcing procedure uses a carry produced by the model’s own rollout, while stopping gradients through that carry. They also report that increasing the number of sampling steps of a standard FLM does not improve Sudoku performance, consistent with our ProsQA observation in Section 3.2. FRM reports a peak accuracy of 99.5%99.5\% on a fixed 10001000-puzzle subset of the Sudoku-Extreme test set, included for context in Table 4. Its maze evaluation uses Maze-Unique rather than Maze-Hard. Looped Flows (Suleymanzade et al., 2026), concurrent with this work, follows the TRM preprocessing and evaluation protocol on Sudoku-Extreme and Maze-Hard. It trains recurrent flow updates with local denoising losses and stops gradients between updates. UFM instead applies one terminal loss and backpropagates through the final NbackN_back updates of the rollout. 7 Limitations Our experiments use relatively small models, from 88M to 1616M parameters, on synthetic structured reasoning tasks. Training through a rollout requires memory for the backpropagated suffix, which limited model size in this work. Extending UFM to larger models and less structured language-reasoning tasks will require more memory-efficient ways to train consecutive updates jointly. The construction in Section 3.1 is an existence result. It relies on orthonormal embeddings and a linear propagation head, and does not show that gradient descent recovers this velocity field. The trained models share only a qualitative pattern with the construction, examined in Section 4.1. Finally, a rollout begins from Gaussian noise, so a correct answer may appear among several trajectories without being selected by a single rollout. Our parameter-free selection rule recovers most of the Pass@100 gap on Sudoku-Extreme, but is less reliable for the much longer Maze-Hard outputs and could benefit from a learned correctness or halting head, as used in HRM and TRM (Wang et al., 2025; Jolicoeur-Martineau, 2025). Although UFM substantially improves over standard flow baselines, it remains behind the strongest reported recursive reasoners on Sudoku-Extreme. 8 Conclusion We show how flow models can use additional integration steps to improve reasoning. On graph reachability, a two-layer construction propagates information along the graph, and its Euler discretization requires at least as many steps as the distance from the root to the answer. Standard flow language models do not use their steps in this way: their accuracy remains largely flat as the number of sampling steps increases. UFM changes this by training through the model’s own latent rollout with a single loss at its endpoint. Accuracy then grows with the number of steps on ProsQA. For longer Sudoku and Maze rollouts, sphere retraction stabilizes the updates and allows an 8.48.4M-parameter model to outperform 28.628.6M standard-flow baselines. Multiple rollouts further improve reliability, although selecting among them is less effective for long answers. Applying this approach to larger language models will require more memory-efficient rollout training and stronger selection mechanisms. Our results show that training through the rollout can make additional flow steps useful for reasoning. Acknowledgements This work was supported in part by the Hector Foundation, by the Max Planck ETH Center for Learning Systems, and by EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1. References Albergo et al. (2025) M. Albergo, N. M. Boffi, and E. Vanden-Eijnden Stochastic interpolants: a unifying framework for flows and diffusions. Journal of Machine Learning Research 26 (209), p. 1–80. External Links: Link Cited by: §1, §2.1. Baek et al. (2026) J. Baek, M. Jo, M. Kim, M. Ren, Y. Bengio, and S. Ahn Generative recursive reasoning. External Links: 2605.19376, Link Cited by: §D.1, Table 5, §6. Clark et al. (2024) K. Clark, P. Vicol, K. Swersky, and D. J. Fleet Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §5.2, §6. Davis et al. (2026) O. Davis, A. Filippova, P. Ablin, V. Turrisi, A. Shidani, M. Cuturi, and L. Béthune Scaling categorical flow maps. External Links: 2605.07820, Link Cited by: §1. Deschenaux and Gulcehre (2026) J. Deschenaux and C. Gulcehre Language modeling with hyperspherical flows. External Links: 2605.11125, Link Cited by: §D.2, §D.4, Table 4, Table 6, §2.2, §5.1, §5.3, Table 1. Eijkelboom et al. (2024) F. Eijkelboom, G. Bartosh, C. A. Naesseth, M. Welling, and J. van de Meent Variational flow matching for graph generation. In Advances in Neural Information Processing Systems, A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (Eds.), Vol. 37, p. 11735–11764. External Links: Document, Link Cited by: §1. Fein-Ashley and Rashidinejad (2026) J. Fein-Ashley and P. Rashidinejad Solve the loop: attractor models for language and reasoning. External Links: 2605.12466, Link Cited by: Table 4, Table 4, §6. Gao et al. (2025) Z. Gao, L. Chen, Y. Xiao, H. Xing, R. Tao, H. Luo, J. Zhou, and B. Dai Universal reasoning model. External Links: 2512.14693, Link Cited by: Table 4, §6. Hao et al. (2026) S. Hao, S. Sukhbaatar, D. Su, X. Li, Z. Hu, J. Weston, and Y. Tian Training large language models to reason in a continuous latent space. External Links: 2412.06769, Link Cited by: §1, §3.2, §6. Helbling et al. (2026) A. Helbling, A. Bryutkin, M. Martino, D. H. Chau, N. Dehmamy, and H. Strobelt Flow reasoning models: turning flows into efficient recurrent reasoners. External Links: 2606.29150, Link Cited by: Table 4, §6. Hu et al. (2026) K. Hu, L. Qiu, Y. Lu, H. Zhao, T. Li, Y. Kim, J. Andreas, and K. He ELF: embedded language flows. External Links: 2605.10938, Link Cited by: §2.2. Huang et al. (2026) B. Huang, Z. Geng, and Z. Kolter Equilibrium reasoners: learning attractors enables scalable reasoning. External Links: 2605.21488, Link Cited by: Table 4, Table 4, Table 5, §6. Jolicoeur-Martineau (2025) A. Jolicoeur-Martineau Less is more: recursive reasoning with tiny networks. External Links: 2510.04871, Link Cited by: §D.1, Table 4, §1, §5.1, §5.1, §5.2, §6, §7. Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret Transformers are RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 119, p. 5156–5165. External Links: Link Cited by: Definition 9. Lee et al. (2026) C. Lee, J. Yoo, M. Agarwal, S. Shah, J. Huang, A. Raghunathan, S. Hong, N. M. Boffi, and J. Kim Flow map language models: one-step language modeling via continuous denoising. External Links: 2602.16813, Link Cited by: Table 4, §1, §2.2, Table 1. Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §1, §2.1. Liu et al. (2023) X. Liu, C. Gong, and Q. Liu Flow straight and fast: learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §1, §2.1. Movahedi et al. (2026) S. Movahedi, V. Milovanović, S. L. Feigin, A. Theus, T. Hofmann, V. Boeva, T. K. Rusch, and A. Orvieto Fixed-point reasoners: stable and adaptive deep looped transformers. External Links: 2606.18206, Link Cited by: Table 4, Table 4, §6. Peebles and Xie (2023) W. Peebles and S. Xie Scalable diffusion models with transformers. In 2023 IEEE/CVF International Conference on Computer Vision (ICCV), Vol. , p. 4172–4182. External Links: Document Cited by: §2.1. Prieto et al. (2025) L. Prieto, M. Barsbey, P. A. M. Mediano, and T. Birdal Grokking at the edge of numerical stability. External Links: 2501.04697, Link Cited by: §4. Roos et al. (2026) D. Roos, O. Davis, F. Eijkelboom, M. Bronstein, M. Welling, İ. İ. Ceylan, L. Ambrogioni, and J. van de Meent Categorical flow maps. External Links: 2602.12233, Link Cited by: §1. Sghaier et al. (2026) A. Sghaier, A. Parviz, and A. Jolicoeur-Martineau Probabilistic tiny recursive model. External Links: 2605.19943, Link Cited by: Table 5, §6. Snell et al. (2024) C. Snell, J. Lee, K. Xu, and A. Kumar Scaling LLM test-time compute optimally can be more effective than scaling model parameters. External Links: 2408.03314, Link Cited by: §1. Suleymanzade et al. (2026) A. Suleymanzade, C. Lee, F. Eijkelboom, N. M. Boffi, İ. İ. Ceylan, and J. Kim Thinking with looped flows. External Links: 2609.11801, Link Cited by: §6. Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin Attention is all you need. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §2.1. Wang et al. (2025) G. Wang, J. Li, Y. Sun, X. Chen, C. Liu, Y. Wu, M. Lu, S. Song, and Y. A. Yadkori Hierarchical reasoning model. External Links: 2506.21734, Link Cited by: §D.1, Table 4, §1, §5.1, §5.1, §6, §7. Wei et al. (2022) J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. H. Chi, Q. V. Le, and D. Zhou Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (Eds.), External Links: Link Cited by: §1. Yoo et al. (2026) J. Yoo, W. Kim, F. Eijkelboom, C. Lee, N. M. Boffi, S. Hong, and J. Kim Self-conditioned flow map language models via fixed-point flows. External Links: 2607.00714, Link Cited by: §1. Zhu et al. (2025) H. Zhu, S. Hao, Z. Hu, J. Jiao, S. Russell, and Y. Tian Reasoning by superposition: a theoretical perspective on chain of continuous thought. External Links: 2505.12514, Link Cited by: §A.1, §A.3, §A.4, §A.5, §A.5, Table 2, §D.2, §1, §3.1, §3.1, §3.1, §3.2, §6, Lemma 5, Lemma 6. Appendix Overview The appendix provides theoretical proofs, training and sampling algorithms, and additional experimental details and results. • Appendix A gives the graph-reachability construction and proofs, including Euler discretization, noisy initialization, and full-vocabulary readout. • Appendix B describes sphere retraction and reports the retraction ablation and norm diagnostic. • Appendix C provides the training and inference algorithms for long rollouts. • Appendix D gives the experimental settings, evaluation procedures, and additional results. Appendix A Graph Reachability: Construction and Proofs We give the construction and proofs supporting Section 3.1. We describe the prompt representation and attention layers, then establish the exact-flow solution, its Euler discretization, and the noisy-initialization bound. A.1 Notation and assumptions Notation is collected in Table 2. Fix a maximum graph size nmax≥2n_ ≥ 2 and a vocabulary containing the vertex labels and the special tokens. These special tokens are distinct from all vertex labels. Graphs are directed, may contain cycles, and have distinct edges, so m≤nmax2m≤ n_ ^2. Each token has an embedding ϕv∈ℝdte φ_v ^d_te. As in Zhu et al. (2025), the construction assumes ⊤= = I, which requires dte≥Vd_te≥ V. These assumptions apply to the construction; the trained models use learned embeddings. We set dist(r,v)=∞dist(r,v)=∞ for unreachable vertices. The walk count includes zero-length walks: W0(u,v)=1W_0(u,v)=1 when u=vu=v and 00 otherwise. The evolving state t z_t is the content slot of ⟨A⟩ ; the prompt information remains in the cache. Table 2: Notation for Section 3.1 and this appendix. The last column gives the symbol of Zhu et al. (2025) where one exists; we rename their tjt_j and u~v u_v to free t for flow time and u,vu,v for vertices. Symbol Meaning Zhu et al. =(,ℰ)G=(V,E) directed graph, n=||≤nmaxn=|V|≤ n_ , m=|ℰ|≥0m=|E|≥ 0 distinct edges same (uj→vj)(u_j\!→\!v_j) source and target of edge j (sj,tj)(s_j,t_j) r; c1,c2c_1,c_2; ci⋆c_i , c3−i⋆c_3-i root; candidates; reachable, unreachable candidate same dist(⋅,⋅)dist(·,·); D=ecc(r)D=ecc(r) graph distance; eccentricity of the root — kV_k vertices within k hops of r cV_c Wk(u,v)W_k(u,v) number of directed walks of length k from u to v — T=3m+7T=3m+7; Tmax=3nmax2+7T_ =3n_ ^2+7 prompt length; its bound — ⟨s⟩,⟨e⟩,⟨Q⟩,⟨R⟩,⟨A⟩ , , , , special tokens same ϕv∈ℝdte φ_v ^d_te, token embedding, ⊤= = I u~v u_v ct,b1,b2,b3,b4,pect,b_1,b_2,b_3,b_4,pe residual-stream slots content, buffers, pos ¯a p_a; εpe _pe sinusoidal encoding of position a; separation gap same; — Pos(⋅)Pos(·) position in the prompt same Ch(x,ℓ)Ch(x, ); κ,βκ,β attention chooser (token x, offset ℓ ); its gains — εth _th threshold — t∈ℝdte z_t ^d_te content slot of ⟨A⟩ at flow time t — A_G; λ propagation operator; value gain — Γu,v _u,v; δ ∑kλkWk(u,v)/k! _kλ^kW_k(u,v)/k!; readout margin — N number of Euler steps C (thoughts) ∼(0,) ξ (0, I); σ initialization noise; its scale — M0M_0 readout gain on the candidates — A.2 Prompt format <s>u1u_1v1v_1<e>u2u_2v2v_2<e>…umu_mvmv_m<e><Q>c1c_1c2c_2<R>rrstart tokendirected edgespecial edge tokensource nodetarget nodequestion tokencandidate nodesreasoning tokenroot node<A>answer token Figure 4: Visualization of the prompt token sequence. A.3 Overview Relation to Zhu et al. (2025). The first layer follows their construction: five attention choosers (Lemma 6) cache the information needed later, and a threshold MLP removes attention leakage. We adapt their construction as follows. (i) As in their paper, we state the construction with causal attention; Remark 7 shows that the same weights work with the bidirectional attention of our models. (i) The last three heads read from ⟨A⟩ : two cache the candidates there, and the fifth writes the root r into the content slot of ⟨A⟩ , which sets the initial state of the flow. (i) Their second layer, which expands the frontier by one hop per continuous thought with softmax attention and a threshold MLP, is replaced by a single linear attention head whose output defines the velocity field. What is evaluated along the flow. Layer 1 is applied once to the prompt: it fills the buffers and sets the initial state 0 z_0 at ⟨A⟩ . At every time t, the velocity is the output of the propagation head with the current state t z_t as query. Neither the chooser heads nor the threshold MLP act on t z_t, and the construction uses no normalization layers, so the velocity is linear in the state. A.4 Transformer preliminaries Prompt length. With ⟨A⟩ appended, the prompt has length T=3m+7T=3m+7. Since the edges are distinct, m≤nmax2m≤ n_ ^2, so T≤Tmax:=3nmax2+7T≤ T_ :=3n_ ^2+7. All gains below are chosen from TmaxT_ , so they do not depend on the graph. Residual stream. We split the residual stream of width d=5dte+dped=5d_te+d_pe into six slots: a content slot ctct and four buffers b1,b2,b3,b4b_1,b_2,b_3,b_4, all of width dted_te, and a positional slot pepe of width dped_pe. Token embeddings live in the content slot and are orthonormal (Section A.1): ⊤= = I, where =[ϕ1,…,ϕV] =[ φ_1,…, φ_V] and V=||V=|V|. In full-stream expressions, token embeddings are zero-padded outside the content slot. We keep positional encodings separate as i p_i; the displayed cached states omit this positional block. Attention. One head of parameters (Q,K,V,O)( W_Q, W_K, W_V, W_O) computes the query qi=Q(i+i)q_i= W_Q( h_i+ p_i), the keys kj=K(j+j)k_j= W_K( h_j+ p_j), and the attention scores si,j=exp⟨qi,kj⟩∑j′exp⟨qi,kj′⟩,s_i,j= q_i,k_j _j q_i,k_j , where j′j ranges over j′≤ij ≤ i with causal masking and over all of [T][T] in the bidirectional case (Remark 7). Definition 4 (Sinusoidal Positional Encoding). Fix an even positional dimension dpe≥2d_pe≥ 2, set M=Tmax+4M=T_ +4, and let ω=M−2/dpeω=M^-2/d_pe. The positional encoding ¯i∈ℝdpe p_i ^d_pe has entries: p¯i,2k−1=cos(iωk),p¯i,2k=sin(iωk),k∈[dpe/2], p_i,2k-1= (i\,ω^k), p_i,2k= (i\,ω^k), k∈[d_pe/2], and i p_i is the vector living in the positional slot with pe(i)=¯ipe( p_i)= p_i. Lemma 5 (Separation, Zhu et al. 2025). With the choices in Definition 4, ⟨¯i,¯i⟩=dpe/2 p_i, p_i =d_pe/2 for all i, and there exists εpe>0 _pe>0, depending only on TmaxT_ , such that ⟨¯i,¯j⟩≤dpe/2−εpe p_i, p_j ≤ d_pe/2- _pe for all i≠ji≠ j in [Tmax+4][T_ +4]. Hence, for ℓ≤4 ≤ 4 and ℓ<i≤T <i≤ T, ⟨¯i,¯j+ℓ⟩ p_i, p_j+ is maximized over j∈[T]j∈[T] at the unique index j=i−ℓj=i- , with a gap of at least εpe _pe. Lemma 6 (Attention chooser, Zhu et al. 2025). Consider a token x∈x , an offset ℓ∈1,2,3,4 ∈\1,2,3,4\ and ε∈(0,1) ∈(0,1). Under Definition 4, there exist Q,K W_Q, W_K for a causal head, with gains depending only on TmaxT_ and ε , such that for any T≤TmaxT≤ T_ and any input sequence satisfying: (H1) i=∑v∈αi,vϕv h_i= _v _i,v\, φ_v with αi,v≥0 _i,v≥ 0 and ∑vαi,v2=1 _v _i,v^2=1, for all i; (H2) ⟨ϕx,i⟩∈0,1 φ_x, h_i ∈\0,1\ for all i, so each state is either equal to ϕx φ_x or orthogonal to it; (H3) ⟨ϕx,i⟩=0 φ_x, h_i =0 for i≤ℓi≤ , so none of the first ℓ positions holds x. We have, for all i∈[T]i∈[T], si,i−ℓ>1−ε if ⟨ϕx,i⟩=1,si,1>1−ε otherwise.s_i,\,i- >1- \ if φ_x, h_i =1, s_i,1>1- \ otherwise. We write Ch(x,ℓ)Ch(x, ) for such a head. The idea is the following. Since the encoding is sinusoidal, ¯j+ℓ=ℓ¯j p_j+ = R_ \, p_j for a fixed rotation ℓ R_ , so a linear key can contain both ¯j+ℓ p_j+ and ¯j p_j. Taking qi∝(¯i,κ⟨ϕx¯,i⟩¯1)q_i ( p_i,\ κ φ_ x, h_i \, p_1 ) and kj∝(ℓ¯j,¯j)k_j ( R_ p_j,\ p_j ) in two orthogonal subspaces, with ϕx¯=∑v≠xϕv φ_ x= _v≠ x φ_v, gives ⟨qi,kj⟩=β(⟨¯i,¯j+ℓ⟩+κ⟨ϕx¯,i⟩⟨¯1,¯j⟩). q_i,k_j =β ( p_i, p_j+ +κ\, φ_ x, h_i \, p_1, p_j ). If position i holds x, then ⟨ϕx¯,i⟩=0 φ_ x, h_i =0 and Lemma 5 puts the unique maximum at j=i−ℓj=i- (and i>ℓi> by (H3), so i−ℓ≥1i- ≥ 1). Otherwise ⟨ϕx¯,i⟩≥1 φ_ x, h_i ≥ 1 by (H1)–(H2), and for κ large the second term, uniquely maximized at j=1j=1, dominates. A large β turns both gaps into attention weights above 1−ε1- . We refer to Zhu et al. (2025) for the complete proof. Our five heads use x∈⟨e⟩,⟨A⟩x∈\ , \ and offsets ℓ∈1,2,3,4 ∈\1,2,3,4\. Each input position carries one token embedding, so (H1)–(H2) hold. Edge tokens occur at positions 3j+1≥43j+1≥ 4, while the answer token occurs at position 3m+7≥73m+7≥ 7. Since the edge heads use offsets at most two and the answer heads use offsets at most four, (H3) also holds. Remark 7 (Bidirectional attention). In a bidirectional head, the softmax runs over all j∈[T]j∈[T] instead of j≤ij≤ i, and the score is unchanged: ⟨qi,kj⟩=β(⟨¯i,¯j+ℓ⟩+κ⟨ϕx¯,i⟩⟨¯1,¯j⟩). q_i,k_j =β ( p_i, p_j+ +κ\, φ_ x, h_i \, p_1, p_j ). For a head that does not fire, ⟨ϕx¯,i⟩≥1 φ_ x, h_i ≥ 1 by (H1)–(H2), and by Lemma 5, ⟨¯1,¯j⟩ p_1, p_j is uniquely maximized at j=1j=1 over all of [T][T], not only over j≤ij≤ i; for κ large, the head attends to ⟨s⟩ . For a head that fires, ⟨ϕx¯,i⟩=0 φ_ x, h_i =0 and the score is β⟨¯i,¯j+ℓ⟩β p_i, p_j+ ; since j+ℓ≤T+4≤Tmax+4j+ ≤ T+4≤ T_ +4 for every j∈[T]j∈[T], Lemma 5 gives the unique maximum at j=i−ℓj=i- with the same gap εpe _pe. The softmax runs over at most TmaxT_ positions in both cases, so the gains β,κβ,κ of Lemma 6 are unchanged, and Lemma 8 holds. The propagation head is unaffected, since ⟨A⟩ is the last position. A.5 Layer 1: caching Lemma 8 (Static memory and answer initialization). Append ⟨A⟩ immediately after the root. One pass of the five chooser heads Ch(⟨e⟩,2),Ch(⟨e⟩,1),Ch(⟨A⟩,4),Ch(⟨A⟩,3),Ch(⟨A⟩,1),Ch( ,2), ( ,1), ( ,4), ( ,3), ( ,1), followed by a threshold MLP, stores the source and target of each edge at its edge token, stores ϕc1 φ_c_1 and ϕc2 φ_c_2 in separate buffers at ⟨A⟩ , and copies ϕr φ_r into the content slot of ⟨A⟩ . The relevant states are ⟨e⟩,j=[ϕ⟨e⟩ϕujϕvj],⟨A⟩=[ϕ⟨A⟩+ϕrϕc1ϕc2],so that0=ϕ⟨A⟩+ϕr. h_ ,j= bmatrix φ_ \\ φ_u_j\\ φ_v_j\\ 0\\ 0 bmatrix, h_ = bmatrix φ_ + φ_r\\ 0\\ 0\\ φ_c_1\\ φ_c_2 bmatrix, that z_0= φ_ + φ_r. (8) The key and value of the propagation head (Definition 9) read b1b_1 and b2b_2 of ⟨e⟩,j h_ ,j, and b1b_1 is zero at every non-edge position. Non-firing heads assign more than 1−ε1- attention to ⟨s⟩ , whose value is zero. We choose εth=1/(2nmax) _th=1/(2 n_ ) and ε=εth/2 = _th/2. The threshold and chooser gains can therefore be fixed using only nmaxn_ . Proof. Each of the five heads is an attention chooser, so each one puts at least 1−ε1- of its attention on the position it is aimed at. We check that the offsets land correctly. With ⟨A⟩ appended after the root, the positions are Pos(⟨s⟩)=1Pos( )=1, Pos(uj)=3j−1Pos(u_j)=3j-1, Pos(vj)=3jPos(v_j)=3j, Pos(⟨e⟩,j)=3j+1Pos( ,j)=3j+1, Pos(⟨Q⟩)=3m+2Pos( )=3m+2, Pos(c1)=3m+3Pos(c_1)=3m+3, Pos(c2)=3m+4Pos(c_2)=3m+4, Pos(⟨R⟩)=3m+5Pos( )=3m+5, Pos(r)=3m+6Pos(r)=3m+6 and Pos(⟨A⟩)=3m+7Pos( )=3m+7. So Ch(⟨e⟩,2)Ch( ,2) and Ch(⟨e⟩,1)Ch( ,1) read uju_j and vjv_j; Ch(⟨A⟩,4)Ch( ,4) and Ch(⟨A⟩,3)Ch( ,3) read c1c_1 and c2c_2; and Ch(⟨A⟩,1)Ch( ,1) reads r. As in Zhu et al. (2025), all heads use the value matrix, V=[dte−ϕ⟨s⟩ϕ⟨s⟩⊤ 0] W_V=[\, I_d_te- φ_ φ_ \ \ 0\,], which will send ϕ⟨s⟩ φ_ to zero. Heads 1–4 write into buffers 1–4, respectively, and head 5 writes into the content slot. A non-firing head assigns more than 1−ε1- attention to ⟨s⟩ , whose value is zero. Each remaining output coefficient is therefore smaller than ε and is removed by the threshold MLP. As in Zhu et al. (2025), a threshold MLP with threshold εth _th, applied to each coefficient of the content slot and the buffers in the embedding basis, maps coefficients at most εth _th to 00 and coefficients at least 1−εth1- _th to 11. Each head puts attention above 1−ε1- on its target and leaks less than ε elsewhere, so with ε<εth < _th every coefficient becomes exactly 00 or 11, and εth _th does not depend on the graph. Only head 1 writes into buffer 1, and it fires only at edge tokens, so after thresholding buffer 1 is zero at every other position. Since buffer 1 is the key of the propagation head, only edge tokens contribute to its output. Buffer 2 is written only by head 2, which fires at edge tokens, and buffers 3 and 4 only by heads 3 and 4, which fire at ⟨A⟩ . Hence at Pos(⟨A⟩)Pos( ) buffers 1 and 2 are empty, buffers 3 and 4 hold ϕc1 φ_c_1 and ϕc2 φ_c_2, and the content slot holds ϕ⟨A⟩ φ_ together with the root written by head 5, so 0=ϕ⟨A⟩+ϕr z_0= φ_ + φ_r, as claimed. ∎ A.6 Layer 2: linear propagation Definition 9 (Linear propagation head). The second layer is a single unnormalized linear-attention head (Katharopoulos et al., 2020), with Q=[dte 0],K=[ 0dtedte 0],V=λ[ 0dte 0dtedte 0],O=[dte 0]⊤, W_Q=[\, I_d_te\ \ 0\,], W_K=[\, 0_d_te\ \ I_d_te\ \ 0\,], W_V=λ\,[\, 0_d_te\ \ 0_d_te\ \ I_d_te\ \ 0\,], W_O=[\, I_d_te\ \ 0\,] , in a way that it reads the content slot as a query, the buffer 11 as key and the buffer 22 as a value with a gain λ, and writes back on the content slot. There is no softmax and no normalizing denominator: oi=O∑j≤i⟨Qi,Kj⟩Vj.o_i= W_O _j≤ i W_Q h_i,\, W_K h_j \, W_V h_j. (the sum runs over all j in the bidirectional case; at ⟨A⟩ , the last position, the two coincide) Define :=∑(uj→vj)∈ℰϕvjϕuj⊤. A_G:= _(u_j→ v_j) φ_v_j\, φ_u_j . (9) By Lemma 8, the buffer 11 is non-zero only at the edge tokens, where it contains ϕuj φ_u_j and the buffer 22 contains ϕvj φ_v_j. Hence, with a query z at ⟨A⟩ , the head returns λλ A_G\, z. Moreover, the terms of the sum ϕvjϕuj⊤ φ_v_j φ_u_j are orthonormal in the Frobenius sense for distinct edges, thus ∥2≤∥F=m. A_G _2≤ A_G _F= m. A.7 Proof of Theorem 1 Because ϕ⟨A⟩ φ_ is orthogonal to every vertex embedding, ϕ⟨A⟩=, A_G\, φ_ = 0, (10) so under ˙t=λt,0=ϕ⟨A⟩+ϕr,t∈[0,1], z_t=λ A_G z_t, z_0= φ_ + φ_r, t∈[0,1], (11) the answer-token identity remains fixed. Writing t=ϕ⟨A⟩+at z_t= φ_ +a_t (12) is only a mathematical decomposition of the single residual state, not a separate architectural block. Its vertex component satisfies a˙t=λat,a0=ϕr. a_t=λ A_Ga_t, a_0= φ_r. (13) Lemma 10 (Reachability from the matrix exponential). Let Wk(r,v)W_k(r,v) be the number of directed walks of length k from r to v. Then at=eλtϕra_t=e^λ t A_G φ_r and, for every v∈v , (i) ⟨ϕv,t⟩=⟨ϕv,at⟩=∑k≥0(λt)k!Wk(r,v)≥0 φ_v, z_t = φ_v,a_t = _k≥ 0 (λ t)^kk!W_k(r,v)≥ 0; (i) for t>0t>0, this coordinate is positive iff v is reachable from r, and is exactly 00 otherwise; (i) if v is reachable, then ⟨ϕv,1⟩≥δ:=min0≤k≤nmax−1λk!>0 φ_v, z_1 ≥δ:= _0≤ k≤ n_ -1 λ^kk!>0. Proof. Orthonormality gives ϕx=∑j:uj=xϕvj A_G\, φ_x= _j:u_j=x φ_v_j, hence by induction kϕu=∑vWk(u,v)ϕv A_G^k φ_u= _vW_k(u,v) φ_v for every vertex u. Expanding the exponential proves (i). All terms are nonnegative, proving (i). For reachable v, a shortest path has length k∗=dist(r,v)≤nmax−1k^*=dist(r,v)≤ n_ -1 and Wk∗(r,v)≥1W_k^*(r,v)≥ 1, so its single series term proves (i). The equality ⟨ϕv,t⟩=⟨ϕv,at⟩ φ_v, z_t = φ_v,a_t follows from ϕv⟂ϕ⟨A⟩ φ_v φ_ . ∎ Proof of Theorem 1. By Lemma 8, layer 1 sets 0=ϕ⟨A⟩+ϕr z_0= φ_ + φ_r. By Definition 9, the velocity at ⟨A⟩ is λtλ A_G z_t, whose time-11 flow is 1=ϕ⟨A⟩+eλϕr z_1= φ_ +e^λ A_G φ_r since ϕ⟨A⟩= A_G φ_ = 0. Lemma 10(i)–(i) give the coefficients and their signs; the reachable candidate has coefficient at least δ by (i) and the unreachable one exactly 00 by (i), which gives the argmax and the margin. The weights depend only on nmaxn_ , the fixed vocabulary, and λ, as required. ∎ A.8 Euler discretization Proof of Corollary 2. Since ϕ⟨A⟩= A_G φ_ = 0, Euler integration gives tN=ϕ⟨A⟩+∏i=0N−1(+λΔti)ϕr,Δti=ti+1−ti>0. z_t_N= φ_ + _i=0^N-1( I+λ t_i A_G) φ_r, t_i=t_i+1-t_i>0. Expanding the product gives ∑j=0Najj _j=0^Na_j A_G^j, where a0=1a_0=1 and every aj>0a_j>0: each coefficient is a sum of products of positive step sizes and powers of λ. By Lemma 10, ⟨ϕv,tN⟩=∑j=0NajWj(r,v). φ_v, z_t_N = _j=0^Na_jW_j(r,v). This coefficient is positive exactly when a walk of length at most N connects r to v, equivalently when dist(r,v)≤Ndist(r,v)≤ N. The reachable candidate therefore has a strictly larger score once N≥dist(r,ci⋆)N (r,c_i ). If N<dist(r,ci⋆)N<dist(r,c_i ), both candidate scores are zero, so the readout does not yet separate them. ∎ A.9 Noisy initialization Proof of Remark 3. For the exact flow with 0=ϕ⟨A⟩+ϕr+σ z_0= φ_ + φ_r+σ ξ, where σ>0σ>0 and ∼(0,) ξ (0, I), t=ϕ⟨A⟩+eλtϕr+σeλt. z_t= φ_ +e^λ t A_G φ_r+σ e^λ t A_G ξ. Let gv=⟨ϕv,⟩g_v= φ_v, ξ , which are i.i.d. (0,1)N(0,1) by orthonormality, and Γv,w=∑kλkWk(v,w)/k!≥0 _v,w= _kλ^kW_k(v,w)/k!≥ 0. By Lemma 10, for every vertex w, ⟨ϕw,1⟩=Γr,w+σNw,Nw=∑vgvΓv,w∼(0,ρw2), φ_w, z_1 = _r,w+σ N_w, N_w= _vg_v _v,w (0, _w^2), with 1≤ρw≤eλm1≤ _w≤ e^λ m. Since every Wk(v,w)≥0W_k(v,w)≥ 0, the covariance of Nc1N_c_1 and Nc2N_c_2 is ∑vΓv,c1Γv,c2≥0 _v _v,c_1 _v,c_2≥ 0, so Var(Nc3−i⋆−Nci⋆)≤ρc12+ρc22≤2e2λmVar(N_c_3-i -N_c_i )≤ _c_1^2+ _c_2^2≤ 2e^2λ m. The readout fails only if σ(Nc3−i⋆−Nci⋆)≥Γr,ci⋆σ(N_c_3-i -N_c_i )≥ _r,c_i , and Γr,ci⋆≥δ _r,c_i ≥δ by Lemma 10(i). The tail bound Pr[Z≥u]≤e−u2/2 [Z≥ u]≤ e^-u^2/2 for Z∼(0,1)Z (0,1) and u≥0u≥ 0 gives Pr[argmaxv∈c1,c2⟨ϕv,1⟩=ci⋆]≥1−exp(−δ24σ2e2λm). \! [ _v∈\c_1,c_2\ φ_v, z_1 =c_i ]≥ 1- \! (- δ^24σ^2e^2λ m ). The component of ξ orthogonal to the vertex span is frozen and invisible to the readout, and [t]=ϕ⟨A⟩+eλtϕrE[ z_t]= φ_ +e^λ t A_G φ_r. ∎ A.10 Readout over the full vocabulary Remark 11 (Full-vocabulary readout). Under the noiseless initialization of Theorem 1, a reachable vertex other than ci⋆c_i can have a larger coefficient because the coefficients count walks. We use the candidate embeddings stored in buffers 3 and 4 to define Sv=⟨ϕv,1⟩+M0(⟨ϕv,ϕc1⟩+⟨ϕv,ϕc2⟩).v∈.S_v= φ_v, z_1 +M_0 ( φ_v, φ_c_1 + φ_v, φ_c_2 ). v . By orthonormality, the added term is M0M_0 on the two candidates and zero elsewhere. Every vertex coefficient is at most eλnmaxe^λ n_ ; the answer-token coefficient is one, and all other special-token coefficients are zero. Thus, choosing M0=2eλnmaxM_0=2e^λ n_ gives Sci⋆≥M0+δS_c_i ≥ M_0+δ, Sc3−i⋆=M0S_c_3-i =M_0, and Sv<M0S_v<M_0 for every other token. The full-vocabulary readout therefore returns ci⋆c_i with margin at least δ. This is a fixed linear decoder on the content slot and candidate buffers, with M0M_0 depending only on nmaxn_ and λ. The velocity field is unchanged. Appendix B Sphere Retraction B.1 Local geometric effect For a nonzero position vector, define ψ(z)=Rz/∥z∥2ψ(z)=Rz/ z _2, with R=dR= d. The following statement describes the local effect of this retraction. Proposition 12 (Local effect of sphere retraction). Let v(z,t)v(z,t) be continuously differentiable near the sphere ∥z∥2=R z _2=R over a fixed interval [0,T][0,T], with T<1T<1. For z on the sphere and sufficiently small Δt t, ψ(z+Δtv(z,t))=z+ΔtPz^v(z,t)+(Δt2),Pz^=−z^z^⊤,z^=z/R.ψ (z+ t\,v(z,t) )=z+ t\,P_ zv(z,t)+O( t^2), P_ z= I- z z , z=z/R. The associated velocity is tangent to the sphere. Proof. For nonzero z, Dψ(z)=R∥z∥2(−zz⊤∥z∥22).Dψ(z)= R z _2 ( I- z z _2^2 ). Taylor expansion at a point with ∥z∥2=R z _2=R gives the stated formula. Since ⟨z,Pz^v⟩=0 z,P_ zv =0, the resulting continuous-time dynamics preserve the norm. ∎ Relation to the implemented update. The algorithms in Appendix C evaluate the denoiser at the raw state and use its retracted version as the update base. If their difference is (h)O(h), where h is the largest grid spacing, local Lipschitz continuity of the denoiser makes the prediction difference (h)O(h). On an interval bounded away from t=1t=1, this changes an update by (Δtkh)O( t_kh). The proposition therefore describes the leading-order interior dynamics under these conditions. It does not establish convergence through the endpoint t=1t=1. B.2 Retraction ablation and norm diagnostic We trained two Sudoku-Extreme models with and without ψ, using the same time-conditioned, constant-learning-rate recipe. At Neval=64N_eval=64, their best recorded test accuracies were 67%67\% with retraction and 55%55\% without it. Both models were trained separately. The ablation uses a different recipe from the cosine-schedule runs in the main comparison. In the norm diagnostic (Figure 5), evaluated over a 252252-step rollout, the model trained without retraction reaches a latent norm roughly 44,00044,000 times its initial scale. Retraction limits this growth, but does not hold the raw state at a fixed norm. The fixed radius applies to the retracted update base. The construction in Section 3.1 also allows the latent norm to grow during propagation, while its final candidate decision is unchanged by positive rescaling. Figure 5: Raw latent norms in separately trained Sudoku-Extreme models with and without retraction. The dashed line marks 448 448, the norm of the retracted update base. The raw state itself need not remain on this sphere. Appendix C Long-Rollout Training and Inference This appendix gives the full procedure used for the long Sudoku and Maze rollouts. Sphere retraction is applied independently to each answer position: ψ((i))=d(i)∥(i)∥2,i=1,…,L.ψ( z^(i))= d\, z^(i) z^(i) _2, i=1,…,L. The denoiser reads the raw state. After the first step, retraction is applied only to the update base. Noise scales and label smoothing are given in Table 7. Algorithm 2 Long-rollout UFM training on one batch with sphere retraction 0: Batch (c,)(c, y), parameters θ, rollout length NtrainN_train, backpropagation length 1≤Nback≤Ntrain1≤ N_back≤ N_train, time-grid distribution q, noise scale σ, and retraction ψ 1: Sample noise ∼(0,) ξ (0, I) in the latent answer space 2: Sample 0≤t0<t1<⋯<tNtrain=10≤ t_0<t_1<·s<t_N_train=1 from q 3: t0←(1−t0)σ+t0Embed() z_t_0←(1-t_0)σ ξ+t_0Embed( y) 4: for k=0k=0 to Ntrain−1N_train-1 do 5: Evaluate this update without gradient tracking when k<Ntrain−Nbackk<N_train-N_back 6: Δtk←tk+1−tk t_k← t_k+1-t_k 7: ^θ(k)←^θ(tk,tk,c) x_θ^(k)← x_θ( z_t_k,t_k,c) 8: if tk+1=1t_k+1=1 then 9: tk+1←^θ(k) z_t_k+1← x_θ^(k) final endpoint prediction 10: else 11: ytk←tky_t_k← z_t_k 12: if k>0k>0 then 13: ytk←ψ(tk)y_t_k←ψ( z_t_k) 14: end if 15: θ(k)←(^θ(k)−ytk)/(1−tk) v_θ^(k)←( x_θ^(k)-y_t_k)/(1-t_k) 16: tk+1←ytk+Δtkθ(k) z_t_k+1← y_t_k+ t_k\, v_θ^(k) 17: end if 18: if k<Ntrain−Nbackk<N_train-N_back then 19: tk+1←sg(tk+1) z_t_k+1 ( z_t_k+1) 20: end if 21: end for 22: ^←StableMax(RMSNorm(tNtrain)dec⊤) y (RMSNorm( z_t_N_train) W_dec ) decode once 23: ℒ←CE(^,)L ( y, y) with the specified label smoothing 24: return ℒ,∇θℒL, _θL Algorithm 3 UFM inference with sphere retraction 0: Prompt c, EMA parameters θEMAθ^EMA, evaluation grid 0=t0<⋯<tNeval=10=t_0<·s<t_N_eval=1, noise scale σ, and retraction ψ 1: Use θEMAθ^EMA for the denoiser and decoder 2: Sample ∼(0,) ξ (0, I) in the latent answer space 3: t0←σ z_t_0←σ ξ 4: for k=0k=0 to Neval−1N_eval-1 do 5: Δtk←tk+1−tk t_k← t_k+1-t_k 6: ^θ(k)←^θ(tk,tk,c) x_θ^(k)← x_θ( z_t_k,t_k,c) 7: if tk+1=1t_k+1=1 then 8: tk+1←^θ(k) z_t_k+1← x_θ^(k) final endpoint prediction 9: else 10: ytk←tky_t_k← z_t_k 11: if k>0k>0 then 12: ytk←ψ(tk)y_t_k←ψ( z_t_k) 13: end if 14: θ(k)←(^θ(k)−ytk)/(1−tk) v_θ^(k)←( x_θ^(k)-y_t_k)/(1-t_k) 15: tk+1←ytk+Δtkθ(k) z_t_k+1← y_t_k+ t_k\, v_θ^(k) 16: end if 17: end for 18: ^←StableMax(RMSNorm(tNeval)dec⊤) y (RMSNorm( z_t_N_eval) W_dec ) 19: return The highest-probability token at each position Appendix D Experimental Details This appendix describes the data splits, training settings, and evaluation procedures used in our experiments. Reporting convention. For the single-rollout results in Table 3, each noise seed gives one independently initialized rollout per test example, using a fixed checkpoint. We report the mean and standard deviation of exact-match accuracy over five noise seeds. Values are rounded to one decimal place, so a displayed standard deviation of 0.00.0 may be nonzero before rounding. The reported checkpoints have the DiT’s time input disabled. Table 3: UFM single-rollout accuracy with Neval=128N_eval=128 and a uniform time grid. Results are percentages, reported as mean ± standard deviation over five noise seeds. Benchmark Test examples Pass@1 (%) Sudoku-Extreme 422,786422,786 74.4±0.074.4± 0.0 Sudoku-Hard 2,0002,000 86.9±0.986.9± 0.9 Maze-Hard 1,0001,000 89.3±0.589.3± 0.5 D.1 Selecting among rollouts We evaluate K=100K=100 rollouts per problem on a 20002000-puzzle subset of the Sudoku-Extreme test set and on all 10001000 Maze-Hard test mazes. On the Sudoku subset, Pass@1 is 74.2%74.2\%, Pass@100 is 98.65%98.65\%, and margin-based selection reaches 98.6%98.6\%. On Maze-Hard, margin-based selection reaches 92.0%92.0\%, compared with 95.0%95.0\% Pass@100. We select among the K rollouts using the mean gap between the largest and second-largest raw logits, δ^=1L∑i=1L(ℓi(1)−ℓi(2)). δ= 1L _i=1^L ( _i^(1)- _i^(2) ). Here, ℓi(1) _i^(1) and ℓi(2) _i^(2) are the two largest logits at position i, before probability normalization. This parameter-free score is an output-space analogue of the separation margin in Theorem 1. Selection is less effective on Maze-Hard. A Sudoku answer has 8181 positions, whereas a Maze answer contains 900900 cells describing the maze layout and solution path. Averaging scores over cells shared by different candidate answers can dilute differences in their paths. On a probe of 300300 Maze boards with 5050 rollouts each at Neval=128N_eval=128, a subset of the boards in the evaluation above, 1818 boards are recoverable, meaning that the first rollout is wrong but at least one of the 5050 is correct; the remaining 282282 boards are either already solved by the first rollout or solved by none of the 5050. On these 1818 boards, correct and incorrect rollouts differ in mean margin by 0.070.07, for mean margins around 200200 in logit units, and the chooser picks a correct rollout on 99 of 1818, against 38%38\% for a uniformly random pick among the distinct answers. These diagnostics suggest that the mean logit margin is less reliable for selecting long maze answers and a learned scorer, such as the latent reward model of GRAM (Baek et al., 2026), or the halting mechanism used in Jolicoeur-Martineau (2025); Wang et al. (2025) is a natural replacement. D.2 ProsQA FLM and UFM use a DiT backbone; S-FLM keeps the architecture of Deschenaux and Gulcehre (2026), which was designed for that method. All models have two layers and approximately 15.715.7M parameters and are trained on the same data. As in Zhu et al. (2025), we use the ProsQA questions whose answer is at most four hops from the root: 14,78514,785 training, 257257 validation, and 419419 test questions, out of the 17,88617,886 / 300300 / 500500 released. UFM is trained with Ntrain=5N_train=5 steps, so each example costs it 55 forward and backward passes. To approximately match the number of forward and backward passes, FLM and S-FLM see 5×5× more training examples: 43,50043,500 optimizer steps at batch size 512512 (≈22.3≈22.3M examples), against 34,50034,500 steps at batch size 128128 for UFM (≈4.4≈4.4M examples). UFM uses a smaller batch to fit the rollout in memory. Gradients flow through all five updates, but the gradient passed back through the latent state is scaled by 0.70.7 at each update; the forward value is unchanged. D.3 Comparison with Recursive Reasoners We also place UFM in context against published recursive reasoners. These rows are reported from the respective papers rather than reruns, so they provide a reference point rather than a controlled pairwise comparison. Table 4: Test accuracy (%) on Sudoku-Extreme and Maze-Hard with a single inference trajectory (Pass@1). Recursive-reasoner numbers are as reported in the respective papers; we did not rerun them. FLM and S-FLM are our runs of the authors’ released code with the default greedy final step, at their best checkpoint and inference setting. UFM uses Neval=128N_eval=128 on a uniform grid; its results are means and standard deviations over five noise seeds, evaluated on 422,786422,786 Sudoku-Extreme puzzles and 10001000 Maze-Hard mazes. FRM’s entry is its reported peak on a fixed 10001000-puzzle Sudoku-Extreme test subset. Its maze evaluation uses Maze-Unique, so no Maze-Hard result is listed. Model Params Sudoku-Extreme Maze-Hard Recursive reasoners Attractor Model (Fein-Ashley and Rashidinejad, 2026) 7M 54.3% 46.7% HRM (Wang et al., 2025) 27M 55.0% 74.5% URM (Gao et al., 2025) 14M 77.6% – TRMa (Jolicoeur-Martineau, 2025) 5M/7M 84.1% 85.3% Attractor Modelb (Fein-Ashley and Rashidinejad, 2026) 27M 91.4% 93.1% EqR (D=64D=64)c (Huang et al., 2026) 5M 93.0% – FPRM (Movahedi et al., 2026) 7M 94.2% 87.0% Flow models FLM (Lee et al., 2026) 28.6M† 10.7% 40.4% S-FLM (Deschenaux and Gulcehre, 2026) 28.6M† 09.4% 49.9% UFM (ours) 8.4M 74.4±0.0%74.4± 0.0\% 89.3±0.5%89.3± 0.5\% FRM (Helbling et al., 2026) 7M 99.5%∗ – aTRM-MLP (5M) on Sudoku, reproduced value from Huang et al. (2026); the original paper reports 87.4%. TRM-Att (7M) on Maze-Hard as reported. bNot reproduced by Movahedi et al. (2026), who obtain 71.4% on Sudoku. cD=64D=64 outer iterations at test time, 4×4× the training budget; at D=16D=16 the result is 86.4%. †At 88M, both baselines failed to train stably, reaching 1.5%1.5\% on Sudoku-Extreme and then declining. ∗FRM reports this peak on a fixed 10001000-puzzle subset of the Sudoku-Extreme test set, not the full 422,786422,786-puzzle test set used for our Pass@1 evaluation. Table 5: Test accuracy (%) on Sudoku-Extreme and Maze-Hard with multiple inference rollouts. Rollouts is the number of trajectories drawn at inference, followed by selection or voting. Recursive-reasoner numbers are as reported in the respective papers. UFM uses Neval=128N_eval=128 Euler steps and K=100K=100 rollouts with the chooser of Section 5.4. UFM is evaluated on a 20002000-puzzle Sudoku-Extreme test subset and all 10001000 Maze-Hard test mazes (Section D.1). Model Params Rollouts Sudoku-Extreme Maze-Hard Recursive reasoners GRAMa (Baek et al., 2026) 10M 20 97.0% – PTRMb (Sghaier et al., 2026) 5M/7M 100 98.75% 86.73% EqR (D=64D=64)c,d (Huang et al., 2026) 5M 128 99.8% – Flow models UFM (ours) 8.4M 100 98.6% 92.0% a2020 parallel trajectories selected by a learned latent reward model. bInference-only on a TRM checkpoint; best-Q@100 on Sudoku, majority vote on Maze; Pass@100 is 99.06% / 95.63%. cD=64D=64 outer iterations at test time, 4×4× the training budget. dEqR reports 93.0% on Maze-Unique, a filtered variant, not Maze-Hard. D.4 Baseline replication on Sudoku-Hard We use the best tested checkpoint of each baseline at 3030k training steps, corresponding to 1.61.6 hours for FLM and 1.71.7 hours for S-FLM. For S-FLM, we used the truncated schedule in the authors’ released code (the 43.2%43.2\% setting in Table 1 of Deschenaux and Gulcehre (2026)), not the adaptive-schedule variant. The released split has no validation set, so each baseline checkpoint is selected on the same 22k puzzles on which it is reported, which favors the baselines. The UFM checkpoint is instead selected on a separate 22k validation split (best exact match at N=64N=64); this is the final checkpoint (19.719.7k steps, 1.11.1 h). The updated evaluation reaches 86.9±0.9%86.9± 0.9\% at N=128N=128 over five noise seeds. Table 6 reports accuracy as a function of the number of sampling steps, with the default greedy final step. With ancestral noise removal at N=128N=128, S-FLM reaches 50.9%50.9\%, the number in Table 1; FLM has no such option. Training times exclude evaluation. Table 6: Accuracy (%) on the hard Sudoku split of Deschenaux and Gulcehre (2026) as a function of the number of sampling steps N, with the default greedy final step and without ancestral noise removal. FLM and S-FLM are our runs of the authors’ released code (28.628.6M parameters) with the cosine schedule of their code base, at their best checkpoint (3030k steps), selected on these same puzzles, which favors the baselines. With ancestral noise removal at N=128N=128, S-FLM reaches 50.9%50.9\% (Table 1); with the released constant schedule it reaches 36.8%36.8\%, and its authors report 43.2%43.2\% with the truncated linear noise schedule and 45.0%45.0\% with the additional adaptive schedule. UFM (8.48.4M parameters) is trained with 2424-step rollouts and selected on a separate validation split. Training is training-only wall-clock time on one A100. The UFM row reports an earlier single-run step-count sweep. The updated five-noise-seed result at N=128N=128 is 86.9±0.9%86.9± 0.9\% (Table 3). The second UFM row uses the same setup with a longer training budget of 2222 hours. We report it separately from the main comparison in Table 1. Model Train (h) N=8N=8 1616 3232 6464 128128 FLM 1.6 h 51.2 51.8 51.8 51.6 51.9 S-FLM 1.7 h 48.9 49.5 50.0 50.2 50.1 UFM (ours) 1.1 h 2.8 65.1 80.2 84.4 86.1 UFM (ours) 22 h 64.1 97.0 99.4 99.8 99.8 D.5 Hyperparameters UFM on Sudoku and Maze. All three benchmarks use a two-layer DiT with hidden width 448448, eight attention heads, and 8.48.4M parameters. The backbone uses pre-norm RMSNorm, SwiGLU with expansion factor four, adaLN-Zero, two-dimensional rotary position embeddings, and 1616 puzzle prefix tokens. The implementation supports sinusoidal time conditioning with frequency dimension 256256. The reported checkpoints disable this input; flow time still enters through the update coefficients. We use AdamW with β=(0.9,0.95)β=(0.9,0.95), gradient clipping at 1.01.0, dropout 0.10.1, and an exponential moving average of the weights with decay 0.9990.999. After warmup, the learning rate follows a cosine schedule from 10−410^-4 to 10−510^-5. We retain weights and accumulation in FP32 and use BF16 matrix multiplications. Settings that differ across benchmarks are listed in Table 7. Table 7: Benchmark-specific UFM training settings. The noise scale σ is the standard deviation of each latent noise coordinate. The Sudoku-Hard column describes the run used for the main comparison. Setting Sudoku-Extreme Maze-Hard Sudoku-Hard Global batch size 768768 128128 256256 Noise scale σ 1/4481/ 448 11 1/4481/ 448 Label smoothing 0.10.1 00 0.10.1 Weight decay 0.750.75 1.01.0 0.750.75 Warmup steps 2,0002,000 2,0002,000 500500 Training budget (steps) 130130k 125125k 19.719.7k Training grid and terminal loss. We train with Ntrain=24N_train=24 updates and backpropagate through the final Nback=6N_back=6; the preceding updates run without gradient tracking. For each example, we sample t0∼[0,t0,max]t_0 [0,t_0, ] and initialize the state from the noise–answer interpolant. The upper bound t0,maxt_0, decreases from 0.60.6 to 0.20.2 over the first 30%30\% of training. We sample and sort 2323 further times in (t0,1)(t_0,1) independently for each example. The final network call returns the predicted endpoint directly, equivalent to taking the last Euler update to t=1t=1. The decoder applies RMS normalization followed by an untied linear vocabulary projection. We compute one StableMax cross-entropy loss at this endpoint, with the label smoothing shown in Table 7. D.6 Computing Resources UFM training uses one NVIDIA A100 on Sudoku-Extreme and Sudoku-Hard, and two A100s on Maze-Hard. Training time excludes evaluation and is estimated from the median steady-state iteration rate in the logs. On Sudoku-Hard, UFM takes 1.141.14 hours for 19.719.7k optimization steps. FLM and S-FLM take 1.071.07 and 1.121.12 hours, respectively, for 2020k steps. Their selected 3030k-step checkpoints require 1.611.61 and 1.681.68 hours of training. These estimates give the rounded training times reported in the main comparison.