Flow Maps: Learning the Jump

Diffusion models generate images by repeatedly predicting how a noisy sample should change. Each prediction is local; the sampler combines many of them into a trajectory. Flow maps learn the result of an entire time interval, making generation in one or a few steps possible.

The central problem is learning a useful large jump. We will build from noise prediction and velocity fields to maps across finite time intervals, then look at how MeanFlow learns these maps without computing training trajectories. Finally, we will separate improvements to the jump itself from methods that improve the distribution of generated images and videos.

Contents

1. From noise to a velocity field
2. From local steps to a flow map
3. MeanFlow: learning average velocity
4. Making flow map training easier
5. Matching the generated distribution
6. Different training states, posterior samples, and video

Throughout, $x_0$ is data and $x_1$ is noise. Generation runs from $t=1$ toward $t=0$. We define velocity as $v=dx/dt$, so a generation step uses a negative time increment.

1. From noise to a velocity field

Diffusion: learn the direction back to data

Diffusion training starts with a clean example and adds Gaussian noise. A schedule sets the signal amplitude $\alpha_t$ and noise amplitude $\sigma_t$:

\[x_t=\alpha_t x_0+\sigma_t\epsilon, \qquad \epsilon\sim\mathcal N(0,I).\]

A denoiser learns to predict the added noise from $(x_t,t)$. Because multiple noise and data pairs can produce the same noisy input, squared error training learns the conditional mean noise, $\mathbb E[\epsilon\mid x_t]$. This prediction determines the score, $s_t(x)=\nabla_x\log p_t(x)$: a vector pointing toward increasing density at that noise level. Together with the noise schedule, the score defines a probability flow ordinary differential equation (ODE), giving a deterministic route from noise to data. Sampling then becomes a numerical integration problem.

Flow Matching: learn velocity directly

Flow Matching trains the ODE velocity itself. One common construction draws independent data and noise endpoints and connects them with a straight segment:

\[x_t=(1-t)x_0+tx_1,\qquad v_c=x_1-x_0.\]

The sampled segment has constant velocity $v_c$, so training can supervise any intermediate point without running a solver. The figure shows the resulting distributions as time moves from noise toward data.

The network sees only $(x_t,t)$, so it learns $v(x,t)=\mathbb E[v_c\mid x_t=x]$. This is the marginal velocity, averaged over the endpoint pairs that could produce that state. As the sampler moves, the contributing pairs change, and the average direction can turn. Consequently, straight training segments can produce curved sampling trajectories. This is a key reason large sampling steps are difficult.

2. From local steps to a flow map

Why a larger Euler step misses

Euler sampling takes the current velocity and holds it fixed across an interval: $x_s\approx x_t+(s-t)v(x_t,t)$. This works locally, where the direction changes little. Over a longer interval, the true trajectory can bend away from that initial tangent. A more accurate velocity network does not remove this integration error.

For a concrete example, take $\dot x=x$, starting at $x_1=1$. The exact destination is $x_0=e^{-1}\approx0.368$. One Euler step gives $1-1=0$, despite using the exact velocity. More steps reduce the error:

Method	Predicted $x_0$
1 Euler step	0
2 Euler steps	0.250
4 Euler steps	0.316
Exact map: $\Phi(x,t,s)=e^{s-t}x$	0.368

The map predicts a destination

A flow map $\Phi(x_t,t,s)$ returns the state at destination time $s$, given the current state and time. It represents the result of integration:

\[\Phi(x_t,t,s)=x_t+\int_t^s v(x_\tau,\tau)\,d\tau.\]

The teal arrow above is the net displacement between endpoints; the underlying trajectory can still curve. An exact map leaves the state unchanged when $s=t$, and a direct jump agrees with two consecutive jumps through an intermediate time $r$:

\[\begin{gathered} \Phi(x,t,t)=x,\\ \Phi(x,t,s)=\Phi\big(\Phi(x,t,r),r,s\big). \end{gathered}\]

These properties connect different interval lengths and provide constraints for learning a map. In the scalar example, learning $e^{s-t}x$ would let the model return the endpoint directly.

Two related ways to simplify sampling

Rectified Flow makes the trajectories easier to integrate. Its reflow procedure runs the current sampler to pair each noise sample with a generated endpoint, then retrains on straight interpolations between those pairs. Repeating this can straighten transport and improve coarse Euler sampling.

Consistency Models instead learn a shared clean endpoint for states on the same ODE trajectory. During distillation, neighboring teacher states should produce matching predictions; a boundary condition anchors clean inputs to themselves. A standard consistency model fixes the destination at the clean endpoint, while a flow map with two time inputs lets the sampler request an intermediate destination too.

3. MeanFlow: learning average velocity

Replace the integral with an average

MeanFlow parameterizes a jump using its average velocity $u$ over $[r,t]$. Multiplying that average by the interval length gives displacement:

\[(t-r)u(x_t,r,t)=\int_r^t v(x_\tau,\tau)\,d\tau.\]

The rectangle and curved area match even though their heights differ point by point. Once the network predicts $u$, sampling is simply $x_r=x_t-(t-r)u_\theta(x_t,r,t)$.

Learn the average from local supervision

Computing the integral for every training target would require trajectories. MeanFlow avoids this by differentiating the identity with $r$ fixed:

\[u+(t-r)\frac{du}{dt}=v, \qquad \frac{du}{dt}=J_xu\,v+\partial_tu.\]

The derivative changes both the starting state and its time. A Jacobian vector product (JVP) evaluates this combined change without forming the full Jacobian $J_xu$. For inputs $(x,r,t)$, its direction is $(v,0,1)$.

Training and sampling in one step

Training substitutes the sampled conditional velocity $v_c$ and constructs a corrected regression target:

\[\begin{gathered} D_\theta=\operatorname{JVP}[u_\theta;(v_c,0,1)],\\ \mathcal L=\mathbb E\left\|u_\theta-\operatorname{sg}\left[v_c-(t-r)D_\theta\right]\right\|^2. \end{gathered}\]

Here $\operatorname{sg}$ stops gradients through the target, avoiding backpropagation through the JVP. At $r=t$, the correction vanishes and training becomes ordinary velocity regression. Longer intervals teach the average needed for larger jumps. Generation in one step sets $(t,r)=(1,0)$; inference needs one prediction of $u_\theta$, without a JVP. MeanFlow can train from scratch without a teacher.

4. Making flow map training easier

iMF: improve the derivative direction

The JVP multiplies a velocity by the network’s Jacobian. If that velocity has high variance, the Jacobian can amplify it. Improved MeanFlow (iMF) replaces the sampled endpoint difference in the JVP with a predicted marginal velocity. It fits local velocity through a compound prediction consisting of the average velocity plus its derivative correction. This keeps the identity over finite intervals while reducing noise in the derivative direction.

The method also conditions on guidance strength. Classifier free guidance normally combines conditional and unconditional predictions, adding inference work. Conditioning the guided field on a scale allows the student to vary guidance strength while retaining a single prediction per step.

Curricula: establish local motion first

AlphaFlow analyzes a conflict between fitting trajectory velocity and enforcing trajectory consistency. Its objective family connects these constraints, and a curriculum shifts emphasis from velocity fitting toward MeanFlow. Accelerated MeanFlow training reaches a related practical conclusion: establish instantaneous velocity and averages over short intervals before emphasizing long intervals. These methods change the order and weighting of training tasks; they help the network learn large jumps from an already useful local field.

Decoupled MeanFlow: reuse a pretrained representation

A pretrained flow transformer already extracts useful features from a noisy state. Decoupled MeanFlow reuses that computation by conditioning early blocks on current time and final blocks on destination time. The early representation describes the input; the later blocks turn it into the requested jump. Flow Matching pretraining can therefore precede the more expensive flow map training, concentrating fine tuning through derivatives on a model that already understands the data.

TVM: differentiate the other endpoint

MeanFlow changes the start while holding the destination time fixed. Terminal Velocity Matching (TVM) holds the input fixed and changes the destination. For displacement $f(x_t,t,s)=\Phi(x_t,t,s)-x_t$, the derivative with respect to destination time must equal the velocity at the endpoint:

\[\partial_s f(x_t,t,s)=v\big(x_t+f(x_t,t,s),s\big).\]

TVM parameterizes $f_\theta=(s-t)F_\theta$, enforcing zero displacement when the two times match. It trains the terminal derivative against a stopped target from an exponential moving average of the velocity model, with Flow Matching supervision anchoring local velocity. Unlike MeanFlow’s detached JVP target, TVM backpropagates through its derivative prediction. That changes the training memory and autodiff requirements; sampling still only evaluates the learned jump.

5. Matching the generated distribution

DMD: compare where samples accumulate

Learning transport constrains individual jumps. A fast generator can still place too many samples in some regions and too few in others. Distribution Matching Distillation (DMD) supervises this error in the overall distribution. A teacher denoiser estimates the target score $s_{p,t}$, while an auxiliary denoiser learns the current generator score $s_{q,t}$. Both evaluate the same noisy generated sample, so their difference compares the two distributions at the same location and noise level.

The toy student sits to the right of the target. Here $s_{p,t}-s_{q,t}=-2$, so gradient descent moves samples left toward the target. For generated samples $G_\theta(z)$ noised as $x_t=\alpha_tG_\theta(z)+\sigma_t\epsilon$, the gradient of the distribution loss is

\[\nabla_\theta\mathcal L= \mathbb E\left[w(t)\alpha_t J_\theta G_\theta(z)^\top (s_{q,t}-s_{p,t})\right].\]

Here $w(t)$ weights noise levels, and the generator Jacobian translates the score difference into a parameter update. Training alternates fitting the auxiliary denoiser and updating the generator with that difference detached. Original DMD also adds paired regression to stabilize the student. The auxiliary denoiser must keep tracking the generator as its distribution changes.

VSD and SiD: related signals, different objectives

Variational Score Distillation (VSD) applies distribution supervision to generating 3D scenes from text. A differentiable renderer produces images from scene parameters and camera views; an auxiliary denoiser learns the distribution of those rendered images. Comparing it with a pretrained teacher gives a direction that propagates through the renderer into the scene. Accounting for the current rendered distribution replaces the reference based on sampled noise used by standard Score Distillation Sampling (SDS).

Score identity Distillation (SiD) starts from a squared score mismatch objective. It uses denoising identities to express that mismatch through residuals involving generated samples and teacher predictions. SiD also alternates an auxiliary denoiser with a generator, but its generator loss differentiates through score network inputs. The similar setup with two denoisers therefore has different gradient paths from DMD, affecting implementation and training cost.

Adversarial supervision: learn a real/fake signal

A discriminator can learn distribution feedback directly from real and generated examples. Adversarial Diffusion Distillation (ADD) combines discrimination of image features with supervision from a diffusion teacher: one encourages realistic appearance, while the other constrains the student’s predictions. Latent Adversarial Diffusion Distillation (LADD) extracts discriminator features from a pretrained latent diffusion model. Keeping this supervision in latent space avoids repeatedly decoding images and running a separate image feature extractor.

DiffRatio connects classification back to scores. For a classifier $D_t(x)=P(\mathrm{real}\mid x,t)$ trained on balanced real/fake batches, the optimal log odds recover the density ratio:

\[\log\frac{q_t(x)}{p_t(x)}= \log\frac{1-D_t(x)}{D_t(x)}.\]

Taking its spatial gradient gives $s_{q,t}-s_{p,t}$. A classifier conditioned on time can therefore estimate the score difference directly, supplying generator supervision without a pretrained teacher score. Accuracy now depends on the classifier and its input gradients.

6. Different training states, posterior samples, and video

Once a fast sampler works, its surrounding computation matters: obtaining training examples, evaluating possible completions, or repeatedly running a video backbone. The methods below change the training inputs or computation around the jump; long video also requires learning from generated history.

FreeFlow: query the teacher on student states

Distillation data can differ from a teacher’s output distribution, particularly after guidance or fine tuning. FreeFlow avoids that mismatch by sampling from the noise prior and querying the teacher along the student’s predicted path. One constraint teaches the teacher’s generation trajectory; another corrects errors using the noising behavior of generated samples. The training states therefore follow what the student actually produces. Here, data free means no external dataset during distillation; the teacher and auxiliary estimation still require compute.

Meta Flow Maps: sample alternative completions

A deterministic map returns one endpoint for a fixed state. Steering with rewards may need several plausible completions to estimate how promising that state is. Meta Flow Maps learn a conditional sampler for the clean posterior $p(x_0\mid x_t)$. Fresh auxiliary noise $\eta_k$ selects different completions while the observed state and time remain fixed:

\[x_0^{(k)}=H_\theta(\eta_k;x_t,t), \qquad \eta_k\overset{\mathrm{iid}}{\sim}\mathcal N(0,I).\]

This posterior comes from the noising/interpolation process: several clean examples can explain the same intermediate observation. The model learns an auxiliary flow conditioned on $(x_t,t)$. Its samples provide differentiable estimates of conditional rewards, reducing the need to run many full trajectories inside each steering step. This changes the object being learned: the output represents a distribution of compatible clean samples.

TMD: spend video compute on a smaller transition head

Video transformers make every evaluation of the full model expensive. Transition Matching Distillation (TMD) splits the model into a large backbone for feature extraction and a small conditional flow head. At an outer sampling step, the backbone supplies features that the head can reuse for inner transition updates. After learning transitions, distribution matching refines the generated outputs. The split gives a way to improve a transition without repeating the entire backbone for every inner update.

AAPT: train with the history inference will produce

Autoregressive video adds another mismatch: training may supply reference history, while inference must continue from generated history. Small errors can then become the context for later mistakes. Autoregressive Adversarial Post Training (AAPT) adapts and distills a video model for causal generation, then adversarially trains its rollouts. Student forcing feeds generated chunks into subsequent chunks, exposing training to the histories the model creates.

Causal attention and a key/value cache reuse earlier context. The discriminator supplies feedback on generated video segments, while rollout training addresses accumulated error over time. Evaluation consequently needs long sequences: sharp individual frames do not establish consistent objects, motion, or scenes.

Across these methods, report network function evaluations (NFE) and full generation latency together. Guidance branches, transition heads, and latent decoding all consume time. Hold prompts, resolution, and guidance fixed, then compare fidelity and diversity alongside speed; for video, also measure consistency over the full rollout.