Learning the integral of a diffusion model (flow maps)

TL;DR

Sander Dieleman’s longest blog post yet: a unifying tutorial on flow maps, neural networks $F(\mathbf{x}_s, s, t)$ that predict any point on a diffusion path from any other point on the same path. Flow maps generalize denoisers (which only predict an infinitesimal tangent) and consistency models (which only predict the endpoint), and are the substrate underneath shortcut models, MeanFlow, consistency models, and progressive distillation. The post adopts the flow-map-matching taxonomy of Boffi/Albergo/Vanden-Eijnden, organizing every known training recipe under three consistency rules — compositionality, Lagrangian, Eulerian — and discusses train-from-scratch vs distill-from-teacher tradeoffs, JVP/backprop choices, and applications beyond sampling speedup (reward-based learning, steerability).

Key claims

A deterministic diffusion sampler defines a bijection between noise and data: each $\mathbf{x}_T$ maps to a unique $\mathbf{x}_0$ along a non-crossing path, and the denoiser predicts the tangent direction at every point on that path [§“Charting paths from noise to data”].
A flow map $F(\mathbf{x}_s, s, t) = \mathbf{x}_s + \int_s^t v(\mathbf{x}_\tau, \tau)\,\mathrm{d}\tau$ is a global characterization of the same path: it jumps from any source noise level $s$ to any target noise level $t$ in one network evaluation, generalizing consistency models (the $t=0$ special case) [§“Cartography: mapping paths with a flow map”].
A flow map admits two equivalent parameterizations: predict the destination $F(\mathbf{x}_s, s, t)$ , or predict the mean flow / average velocity $V(\mathbf{x}_s, s, t)$ with $F = \mathbf{x}_s + (t-s)V$ ; the limit $s=t$ recovers the instantaneous velocity, so a flow map contains a denoiser inside it [§“Cartography”].
Three consistency rules characterize a valid flow map, and every known training recipe derives from one: compositionality ( $F(F(\mathbf{x}_s,s,t),t,u) = F(\mathbf{x}_s,s,u)$ , used by shortcut/progressive-matching), Lagrangian consistency (the $t$ -derivative equals the velocity, used by MeanFlow-style training), and Eulerian consistency (the $s$ -derivative + Jacobian-vector-product term equals zero) [§“Three notions of consistency”].
Lagrangian and Eulerian consistency are the same rule in different reference frames — derivable as limits of compositionality when one of the two time intervals shrinks to zero — and differ only in whether you track a moving sample (Lagrangian) or fix the target (Eulerian) [§“Constructing loss functions from equalities”].
There is no free lunch: flow maps strictly generalize diffusion models and enable single-step sampling, but every known training recipe bootstraps from the local tangent view that a diffusion model provides — either via an explicitly pretrained teacher, or via a self-distillation EMA copy of the student [§“Cartography”, closing remarks].
Path curvature is the limiting factor for few-step diffusion sampling; flow matching and Reflow attempt to straighten paths, but the more robust answer is to learn the integral itself rather than rely on the integrator [§“Charting paths from noise to data”].

Method

The post is a survey, not an empirical paper, but the central method is the flow-map-matching framework: parameterize a network $F_\theta(\mathbf{x}_s, s, t)$ (two time-step inputs, one extra over a denoiser), and train it to satisfy one of the three consistency rules. The post derives each rule from first principles, then shows how the rule becomes a loss function — typically by squaring the residual of the equality and stopping gradient on one side (so the network learns to match a bootstrapped target). Boffi et al.’s taxonomy is used as the organizing skeleton.

Key choices the post enumerates: (i) to backprop or not to backprop through the bootstrapped target — full backprop gives stronger signal but is unstable; stop-gradient + EMA-teacher is the practical recipe; (ii) train from scratch vs distill from a pretrained diffusion model — from-scratch flow maps require Lagrangian or Eulerian losses (which need a velocity field as supervision) and tend to be harder to optimize than the equivalent distillation setup; (iii) JVP vs finite-difference for the Eulerian gradient term, which is the same engineering choice MeanFlow / shortcut models face. Concretely, the Eulerian residual involves $\nabla_{\mathbf{x}_s} F(\mathbf{x}_s, s, u)\, v(\mathbf{x}_s, s)$ — a Jacobian-vector product the post argues is best computed with forward-mode autodiff (with a pointer to Luma’s open-source jvp_utils).

Results

No new benchmarks; the post catalogs prior empirical results from the literature it surveys (consistency models, MeanFlow, shortcut models, flow-map matching, self-distillation). The contribution is conceptual unification: organizing the 2024–2026 explosion of “few-step diffusion” methods under three consistency rules and two parameterization choices, and clarifying which existing techniques (CM, sCM, MF, shortcut, IMM, flow-map-matching, self-distillation) sit where in that taxonomy.

Why it’s interesting

This is the conceptual map underneath the entire [[diffusion-distillation]] thread on the wiki. Almost every distillation paper filed in the last week — [[2601.09881-transition-matching-distillation-fast-video]] (TMD’s flow-head rollout is a Lagrangian-consistency recipe with multi-step inner rollout), [[2602.09014-arcflow-unleashing-2-step-text]] (ArcFlow’s non-linear analytic ODE integration is a closed-form flow map over the teacher’s PF-ODE), [[2602.03139-diversity-preserved-distribution-matching-distillation]] (DP-DMD’s role-separated student is a flow-map parameterization choice), [[2602.02214-causal-forcing-autoregressive-diffusion-distillation]] (Causal Forcing’s frame-level injectivity argument is exactly the flow-map bijection requirement applied to AR video), [[2605.15141-causal-forcing-scalable-few-step]] (Causal Forcing++ explicitly substitutes causal consistency distillation for offline PF-ODE generation) — fits one cell of this taxonomy. The post also gives the canonical reference for the JVP-vs-finite-difference engineering tradeoff that [[blog-tridao-gram-newton-schulz-fast-hardware-2026-03]] and the broader Muon/Newton-Schulz thread care about for a different reason. Finally, it ends with an explicit pointer to Luma’s open-source jvp_utils — an unusual reverse-direction citation that makes this post directly relevant to internal tooling.