Lateral Intelligence
Technical Report · June 9, 2026

Looped Flow Matching for Few-Step Language Models

Lateral Intelligence Team

We improve continuous diffusion language models in the few-step regime by allowing them to correct previous integration errors, borrowing intuition from recursive reasoning models. On Sudoku hard this increases exact match accuracy from 38% to 72%.

At Lateral Intelligence, we are interested in models that are more flexible than the autoregressive (AR) transformer. AR models suffer an inherent inflexibility from their next-token-prediction structure: they must commit tokens from left to right, which means additional computation can only be spent generating more tokens, rather than deliberating over the generation. It also means AR models cannot self correct or refine their output as they only ever operate on already committed clean tokens.

We are excited about recent advances in iterative-refinement architectures such as continuous diffusion language models and recursive architectures (TRM and HRM), whose designs overcome these shortfalls. These models can scale additional compute under a fixed sequence length and can iteratively refine their own predictions. This mirrors how humans solve difficult problems: reasoning further and correcting previous guesses before committing to an answer, which yields strong performance on reasoning tasks (as we show below). And because they operate over a sequence of latents rather than a next-token framework, these models allow for useful properties such as steerability via classifier-free guidance (which we will show in a future post).

Our fundamental insight is that final generation quality is a function of the entire path of model predictions across sampling steps. Here we improve on continuous diffusion language models by integrating intuitions from the TRM/HRM model class.

TL;DR Flow and diffusion models generate samples by integrating a predicted instantaneous velocity over time. Standard methods assume a fixed schedule of integration step sizes for a first-order Euler approximation, and treat each denoising step as independent of previous predictions. We argue these assumptions are limiting in the few-step regime, where first-order approximation errors are large and the model should account for and correct the compounding errors of its earlier denoising predictions.

Method

We build on the continuous diffusion language model literature, specifically Flow Language Models (FLM). The flow language model defines a flow over one-hot encoded representations of tokens and learns under a denoising cross-entropy loss.

Let the sampled data sequence be \(y \in \mathbb{R}^{L}\) with one-hot encoded representation \(x_1 = f(y) \in \mathbb{R}^{L \times D}\). We construct a randomly sampled interpolant

\[x_t = t\,x_1 + (1-t)\,x_0,\]

where \(x_0 \sim \mathcal{N}(0, I)\) and \(t \in [0,1]\). The flow language model is then trained under the denoising cross-entropy objective

\[L_{\text{FM}} = \mathbb{E}_{t,\,x_t}\!\left[-\sum_l (x_1^l)^{\top} \log \hat{x}_1^l(x_t)\right].\]

To make the model aware of and able to correct its previous predictions, we adopt intuitions from TRM and train with a backpropagation-through-time (BPTT) looped objective. Given a number of sampling timesteps \(T\), we define a set of evenly spaced timepoints \(\{t_i\}_{i=1}^{T}\) with \(t_1 = 0 \to t_T = 1\). We initialize \(x_{t_1} \sim \mathcal{N}(0, I)\) and recursively apply the denoiser, updating the latent state with

\[ \begin{aligned} v_{k} &= \frac{\hat{x}_1(x_{t_k}) - x_{t_k}}{1 - t_k}, \\[4pt] x_{t_{k+1}} &= x_{t_k} + (t_{k+1} - t_{k})\,v_k. \end{aligned} \]

Finally, we apply our loss to the final generation \(x_{t_T}\) and backpropagate through the sampling steps:

\[L_{\text{loop}} = -\sum_l (x_1^l)^{\top} \log \hat{x}_1^l(x_T).\]

Our total objective is a simple linear combination of the two, parameterized by a mixing coefficient \(\gamma\):

\[L_{\text{total}} = \gamma\, L_{\text{loop}} + (1 - \gamma)\, L_{\text{FM}}.\]

Experiment

We establish preliminary results on Sudoku solving, using the same dataset (48,000 train / 2,000 val) as in prior work (Deschenaux, 2026). We adopt from their data-formatting procedure (full details in the Appendix). We evaluate on the hardest puzzle difficulty (Hard: 30/81 clues provided), and adopt the same model hyperparameters as FLM.

Model Objective γ T Correct / Total Accuracy (%)
Discrete-time FLM Flow matching only 0.0 6 670 / 2000 33.50
Flow matching + looped 0.5 1158 / 2000 57.90
Looped loss only 1.0 1109 / 2000 55.45
Flow matching only 0.0 12 503 / 2000 25.15
Flow matching + looped 0.5 1433 / 2000 71.65
Looped loss only 1.0 999 / 2000 49.95
Continuous-time FLM Flow matching only — 12 757 / 2000 37.85
Flow matching only — 256 797 / 2000 39.85
Table 1. Exact-match accuracy on Sudoku (Hard). Discrete-time FLM results are shown for varying loss-mixing coefficient γ; continuous-time FLM uses no γ. Results report the best accuracy over training.

The combination of the original flow-matching loss plus the looped BPTT loss outperforms either objective alone. We believe the flow-matching loss serves as a good anchor — it gives the model structure for what an initial denoising-direction estimate should be — while the BPTT loss lets the model learn to generate a sequence of predictions that offset and refine the approximation errors a flow matching only model would otherwise accumulate.

Ablations

We also consider the looped objective from Looped Transformers, where the cross-entropy loss is applied to every denoising prediction at each timestep \(t_k\), rather than only to the final generation:

\[L_{\text{loop}} = -\sum_{k=1}^{T}\sum_l (x_1^l)^{\top} \log \hat{x}_1^l(x_{t_k}).\]

We find this does not perform as well as simply applying BPTT from the final denoising prediction at \(T = 6\), though it is slightly better at \(T = 12\). IN both cases, it remains worse than the combination of the flow-matching loss and BPTT.

Objective Correct / Total Accuracy (%)
T = 6
Looped Transformer loss (all steps) 761 / 2000 38.05
Looped loss, final step only (γ = 1.0) 1109 / 2000 55.45
T = 12
Looped Transformer loss (all steps) 1133 / 2000 56.65
Looped loss, final step only (γ = 1.0) 999 / 2000 49.95
Table 2. Ablation on Sudoku (Hard) comparing the Looped Transformer loss (denoising loss computed at every sampling step) against the looped loss computed at the final step only (γ = 1.0). Neither includes the flow-matching objective. Results report the best accuracy over training.

Limitations

Our method requires specifying a preset number of sampling timesteps. Unlike typical flow models, this means the number of sampling steps cannot easily be varied at inference time. Moreover, because the looped computation is recursive and sequential, training cost increases and is only tractable for a small number of recursive steps \(T\). Future work should address inference-time adaptive compute and reduce the additional cost incurred during pretraining.

Appendix: model and training details

We adopt the data formatting procedure from (Deschenaux, 2026). We concatentate the initial board (leaving unfilled cells as digit 0) with the board solution. We then have data sequences of length 162, where the model takes in the full sequence as input but the loss is only applied to the solution half of the sequence. No explicit information is given about conditioning token positions; when sampling (both in pre-training and generation) we paste the conditioning initial first board in the first half of the sequence.

All experiments use the FLM codebase with a DiT backbone. We use the following configuration: 8 transformer blocks, hidden width 512, 8 attention heads, and a 128-dimensional conditioning embedding (≈29M parameters). We train with a batch size of 32 or 64 and train until convergence, up to 160k training steps.

Discrete-time looped models are trained and sampled with \(T\) evenly spaced first-order Euler steps; the continuous-time FLM baseline uses the same Euler integrator with \(T \in \{12, 256\}\).

References

  1. Deschenaux, J. et al. Language Modeling with Hyperspherical Flows. arXiv:2605.11125, 2026.
  2. Yang, L., Lee, K., Nowak, R. & Papailiopoulos, D. Looped Transformers are Better at Learning Learning Algorithms. arXiv:2311.12424, 2023.
  3. Wang, G. et al. Hierarchical Reasoning Model. arXiv:2506.21734, 2025. — HRM.
  4. Jolicoeur-Martineau, A. Less is More: Recursive Reasoning with Tiny Networks. arXiv:2510.04871, 2025.
  5. Hu, K., Qiu, L. et al. ELF: Embedded Language Flows. arXiv:2605.10938, 2026.
  6. Tiny Recursive Language Diffusion Models. OpenReview (under review), 2026.
  7. Sapient Intelligence. HRM-Text. Code / model release, 2026.
  8. Baek, J., Jo, M., Kim, M., Ren, M., Bengio, Y. & Ahn, S. Generative Recursive Reasoning. arXiv:2605.19376, 2026.
  9. Asadulaev, A., Banerjee, R., Karray, F. & Takac, M. Your Latent Reasoning is Secretly Policy Improvement Operator. arXiv:2511.16886, 2025.
  10. Lee, S. et al. Looped Diffusion Language Models. arXiv:2605.26106, 2026.
  11. Jo, M., Yoon, J., Deschenaux, J., Gulcehre, C. & Ahn, S. Loopholing Discrete Diffusion: Deterministic Bypass of the Sampling Wall. arXiv:2510.19304, 2025.
  12. Liu, S. et al. Think While You Generate: Discrete Diffusion with Planned Denoising. arXiv:2410.06264, 2024.

At Lateral Intelligence, we are working towards advancing the next generation of language modelling with more natively controllable language models. We're a small team based in San Francisco.

If you're building in this direction and want to join, reach out at join@lateralintelligence.org.

If you're interested in supporting the work, get in touch at invest@lateralintelligence.org.

← All Research