Source-linked AI summary
Flow map matching with stochastic interpolants: A mathematical framework for consistency models
Nicholas M. Boffi, Michael S. Albergo, Eric Vanden-Eijnden
TL;DR
Dynamical generative models can require costly numerical integration, while consistency models lack a comprehensive theoretical framework. FMM learns two-time flow maps through distillation or direct training, and the paper reports high sample quality at reduced computational cost on CIFAR-10 and ImageNet-32.
Problem
Dynamical generative models require inference-time numerical integration, and consistency-model design lacks a comprehensive theoretical framework.
Source: §1 · p. 1Method
FMM learns two-time flow maps through distillation losses from a pre-trained velocity field or direct training without one.
Source: §1 · p. 2Results
Experiments on CIFAR-10 and ImageNet-32 report high sample quality at significantly reduced computational cost, with post-training sampling-step adjustment trading accuracy for efficiency.
Source: §1 · p. 2Takeaways & Limitations
The two-time map provides a way to adjust sampling steps after training to balance accuracy and computational efficiency.
Source: §1 · p. 2The Eulerian distillation bound involves the map's spatial gradient, which may be difficult to control, especially when the learned map is singular or nearly singular.
Source: §3.4 · p. 8
Abstract
from arXiv · showhide
Generative models based on dynamical equations such as flows and diffusions offer exceptional sample quality, but require computationally expensive numerical integration during inference. The advent of consistency models has enabled efficient one-step or few-step generation, yet despite their practical success, a systematic understanding of their design has been hindered by the lack of a comprehensive theoretical framework. Here we introduce Flow Map Matching (FMM), a principled framework for learning the two-time flow map of an underlying dynamical generative model, thereby providing this missing mathematical foundation. Leveraging stochastic interpolants, we propose training objectives both for distillation from a pre-trained velocity field and for direct training of a flow map over an interpolant or a forward diffusion process. Theoretically, we show that FMM unifies and extends a broad class of existing approaches for fast sampling, including consistency models, consistency trajectory models, and progressive distillation. Experiments on CIFAR-10 and ImageNet-32 highlight that our approach can achieve sample quality comparable to flow matching while reducing generation time by a factor of 10-20.
1 Introduction
Dynamical generative models learn transport between base and target distributions, but inference-time integration can be costly. Flow Map Matching (FMM) learns a two-time flow map, providing a framework for distillation and direct training.
1 Introduction
- FMM learns a two-time flow map that transports ODE solutions and supports post-training adjustment of integration steps to balance accuracy and computational efficiency.The map can be distilled from a known velocity field or learned directly, and supports arbitrary base distributions.Source: Figure 1 · p. 2
- The framework characterizes flow maps mathematically, deriving distillation algorithms and connections to consistency models and progressive distillation.The paper connects its framework to learning consistency models [28] [29] [30] and progressive distillation of diffusion models [31] [32].Source: §1 · p. 2
- The Lagrangian distillation loss complements an Eulerian loss, which is the continuous-time limit of consistency distillation [28]; both control teacher–student Wasserstein distance.The authors report that the Lagrangian loss performs well across their experiments.Source: §1 · p. 2
- FMM extends its Lagrangian loss to direct flow-map training without a pre-trained velocity field and adapts the flow-matching signal to capture more global information.The paper also discusses direct training with a modified Eulerian loss and caveats of that procedure.Source: §1 · p. 2
- Experiments on CIFAR-10 and ImageNet-32 report high sample quality at reduced computational cost, with sampling steps adjustable after training to trade accuracy for efficiency.The passage identifies post-training step adjustment as a practical consequence of learning the two-time flow map.Source: §1 · p. 2
2 Related Work
Related work reduces the cost of dynamical generative sampling through path straightening, couplings, and tailored or learned solvers. FMM instead learns a two-time flow map, connecting to consistency and distillation methods while addressing limitations of one-time maps.
Dynamical transport of measure
- Flow matching, rectified flow, and stochastic interpolants connect base and target densities through flexible paths, while diffusion models learn to time-reverse an SDE.The passage situates these approaches within dynamical transport of measure and notes their differing constructions.Source: §2 · pp. 2–3
Reducing simulation costs
- Earlier approaches reduce simulation costs by straightening probability-flow paths, simplifying couplings, or using hand-crafted and learned numerical solvers; FMM instead learns the flow map directly.The paper characterizes direct map learning as avoiding optimal-transport-map estimation and overcoming inherent limitations of numerical integration.Source: §2 · p. 3
Consistency models
Extensions of consistency models
- Consistency trajectory models learn two-time maps for improved multistep sampling, but use adversarial training; generalized and bidirectional variants differ in setting or loss from FMM.The paper presents these existing consistency approaches as special cases of its formalism.Source: §2 · p. 3
Progressive and operator distillation
- Neural operator methods learn a one-time map using generated probability-flow trajectories, while progressive distillation matches several original diffusion steps with fewer samples.The passage says these approaches relate to FMM distillation, although the object distilled differs.Source: §2 · p. 3
3 Flow Map Matching
This section presents Flow Map Matching (FMM), which learns two-time maps for probability-flow ODEs as an alternative to numerical integration. It describes an Eulerian objective and a progressive distillation objective, along with convergence and stability caveats.
3.1 Stochastic interpolants and probability flows
- Stochastic interpolants connect endpoint densities through a process whose marginals match a deterministic probability-flow ODE, whose velocity can be learned by square-loss regression [20].The framework encompasses flow matching and rectified flow, as well as variance-preserving and variance-exploding diffusion constructions.Source: §3.1 · pp. 4–5
3.2 Flow map: definition and characterizations
- The flow map advances ODE solutions by arbitrary time intervals, composes consistently across intermediate times, and generates samples from each intermediate density in one or multiple steps.Its definition assumes a one-sided Lipschitz condition on the drift to ensure existence and uniqueness of ODE solutions.Source: §3.2 · pp. 5–6
3.3 Distillation of a known velocity field
- For known velocity fields, Lagrangian and Eulerian residual losses distill the flow map; the Lagrangian approach is new, while the Eulerian approach recovers continuous-time consistency distillation [28] [29] and consistency trajectory models [30].Evaluating the Lagrangian objective on the interpolant avoids generating a dataset from a pretrained model and permits simulation-free training.Source: §3.3 · pp. 6–7
3.4 Wasserstein control
- Both distillation losses bound the Wasserstein distance between teacher and learned-map densities; the Eulerian bound avoids a Lipschitz constant but involves spatial map gradients, and Lagrangian training performed better in experiments.The authors attribute Eulerian training difficulties to potentially singular or nearly singular learned maps, which can yield undefined gradients, fuzzy checkerboard boundaries, and blurry images.Source: §3.4 · p. 8
3.5 Direct training with flow map matching (FMM)
- Direct FMM learns the flow map without a pretrained velocity field, and symmetric time-pair weighting supports one-step generation or shorter-strip training with multiple steps.The time-pair weights must satisfy w_{t,s} = w_{s,t} to learn both the map and its inverse; strip width determines the number of generation steps.Source: §3.5 · pp. 8–9
3.6 Eulerian estimation or Eulerian distillation?
- The proposed direct Eulerian objective is not equivalent to Eulerian distillation, and the section begins analyzing the difference through its cross term.Source: §3.6 · p. 9
- With the identity boundary condition, the true flow map is a critical point of the Eulerian loss when the gradient uses stop-gradient on the map-dependent velocity term.Without this treatment, that term affects minimization, so the loss minimizer need not match the true flow map.Source: §3.6 · p. 10
- The stop-gradient critical point may be difficult to stabilize because the objective is not guaranteed to decrease, although related approaches have succeeded on large-scale image generation [28].Source: §3.6 · p. 10
3.7 Progressive distillation
- Direct one-step map learning was challenging, but convergence improved significantly with weights restricted to time pairs satisfying |t-s| <= 1/K.The resulting K-step model can be converted to a one-step model with a map-distillation loss similar to progressive distillation [31] and neural operator approaches [32].Source: §3.7 · p. 10
- The unique minimizer of the progressive flow-map objective produces in one step the same output as iterating the teacher map K times.The teacher may be trained and frozen before distillation, or the objective can be used in a single network; staged teacher replacement or an offline dataset can reduce the cost of teacher evaluations.Source: §3.7 · pp. 10–11
4 Numerical Realizations
The experiments compare four flow-map learning approaches on a 2D checkerboard and image-generation tasks, including ImageNet-32 and CIFAR-10.
4.1 2D Illustration
- On the checkerboard, four-step FMM improves on one-step FMM, and PFMM distills the four-step map to one step; the main errors occur at discontinuous checker edges.The one-step FMM map performs as well as the probability-flow ODE reference, while LMD outperforms EMD when distilling the same velocity field.Source: §4.1 · p. 11, Figure 3 · p. 12
4.2 Image Generation
- Image-generation experiments train models from scratch and use FID against real images and Teacher-FID against teacher-generated images to assess distillation.The flow map uses a U-Net with both times embedded for its two-time input.Source: §4.2 · p. 13
- On CIFAR-10, 2- and 4-step LMD and EMD samples far outperform the stochastic interpolant, while LMD and PFMM produce effective samples in 1–4 steps.PFMM relies on flow-map matching objectives without a velocity model; Figure 4 also reports that LMD trains faster and attains lower loss and FID than EMD at fixed training steps.Source: Figure 4 · p. 12, §4.2 · p. 13, Table 1 · p. 13
N DDPM BatchOT FMM (Ours)
- On ImageNet-32, FMM has better few-step FID than DDPM and mini-batch OT, while interpolant methods are marginally better at higher steps.Table 2 reports FID scaling with function evaluations for direct FMM training without distillation.Source: Table 2 · p. 14, §4.2 · p. 14
4.2.1 Efficient Style Transfer
- Consistency style transfer maps a class-conditional image backward to s′=0.3, then forward under a new class, retaining style while replacing the subject.The paper uses this example to illustrate cycle consistency and a potential application of the two-time flow map.Source: §4.2.1 · p. 14, Figure 5 · p. 14
5 Conclusion
The conclusion presents flow-map learning through distillation or direct stochastic-interpolant training. The appendix recalls how probability-flow ODEs transport distributions.
5 Conclusion
- The framework learns generative two-time flow maps by distilling a pre-trained velocity model with Eulerian or Lagrangian losses, or by direct stochastic-interpolant training.The conclusion notes that sample quality improves with more map steps, while substantially fewer steps are needed than for other generative models built on dynamical transport.Source: §5 · p. 15
A.1 Transport equation
- Solutions of the probability-flow ODE push initial conditions to the target distribution, and its transport equation can be used to sample a density from initial conditions.The appendix derives the transport equation through its weak formulation.Source: §A.1 · p. 20
A.2 Stochastic interpolants and probability flows
The appendix connects stochastic interpolants to probability-flow ODEs and recalls how score-based diffusion models yield probability flows that FMM can distill.
A.2 Stochastic interpolants and probability flows
- The stochastic interpolant’s density also describes the probability-flow ODE solution, so its associated flow map can serve as a generative model.The interpolant is formed between endpoint densities using a coupled pair of samples and independent Gaussian noise.Source: §A.2 · pp. 20–21
A.3 Score-based diffusion models
- For variance-preserving score-based diffusion models, the score defines a probability-flow velocity field, and FMM can distill the resulting diffusion model.The appendix also recalls time reversal of the SDE and its associated probability-flow ODE.Source: §A.3 · pp. 21–22
A.4 Connections
The appendix reformulates a score-based diffusion probability flow as a stochastic interpolant. This construction transports Gaussian noise to the data point exactly at the endpoint.
A.4 Connections
- A time-changed score-based diffusion probability flow has a stochastic-interpolant representation that removes the small finite-T bias and maps Gaussian noise to data exactly at t=1.The reformulation transports the Gaussian variable at t=0 to the data point at t=1.Source: §A.4 · p. 23
B Proofs for Section 3
The proofs establish optimization properties for Lagrangian, Eulerian, and direct flow-map objectives, alongside structural results about flow maps. The appendix also frames existing consistency and distillation techniques using two-time flow maps.
B Proofs for Section 3
- The flow map is invertible: composing it with its reverse-time map recovers the input.The proof states this property for all times and inputs in the specified domain.Source: B Proofs for Section 3 · p. 23
- The Lagrangian map-distillation loss has the flow map as its global minimizer under the identity boundary condition.The proof treats the objective as a PINN loss whose integrand vanishes when the flow-map equation holds.Source: B Proofs for Section 3 · p. 23
- The Eulerian map-distillation loss likewise has the flow map as its global minimizer under the identity boundary condition.Its PINN loss is minimized when the Eulerian equation holds.Source: B Proofs for Section 3 · p. 24
- The Lagrangian error-bound argument uses a one-sided Lipschitz condition to bound trajectory spread and permits negative C_t for globally contracting maps.The passage explicitly interprets C_t < 0 as the globally contracting case.Source: B Proofs for Section 3 · p. 24
- The Eulerian estimation objective makes the flow map a critical point, since the flow map zeroes its functional gradient.The proof attributes this to the flow map satisfying the Eulerian equation in that gradient.Source: B Proofs for Section 3 · p. 27
- Progressive flow-map matching has a unique minimizer that produces the same output in one step as the original map iterated K times.The construction uses a sequence of intermediate times and a positive weight function.Source: B Proofs for Section 3 · p. 27
C Relation to existing consistency and distillation techniques
- The appendix recasts consistency models and distillation techniques in two-time flow-map notation to clarify their relation to the framework.The passage introduces this as the purpose of the section on existing techniques.Source: C Relation to existing consistency and… · p. 27
C.1 Flow maps and denoisers
The section distinguishes flow maps from denoisers and relates flow-map objectives to consistency models and neural-operator approaches. It identifies how objectives and sampling setups differ across these connections.
C.1 Flow maps and denoisers
- Replacing the reverse-flow-map input with the interpolant in the objective changes its minimizer to a denoiser, not the flow map.With independent endpoints at s = 0 and t = 1, the one-step denoiser recovers only the target-density mean and is not a consistent generative model.Source: §C.1 · p. 28
C.2 Relation to consistency models
- In the variance-exploding setting, the continuous-time consistency loss is a particular Eulerian distillation loss.The one-time consistency function cannot be iterated: by definition, it maps an observation directly to a target sample.Source: §C.2 · p. 29
C.3 Relation to neural operators
- Neural-operator approaches estimate the flow map using a Fourier Neural Operator and pre-generated trajectories to provide training targets without simulation during learning.The described architecture combines the FNO with a U-Net.Source: §C.3 · p. 30
C.4 Relation to progressive distillation
Progressive distillation repeatedly trains a model to replace two sampler steps with one. The supplied comparison discusses distributional fit and preservation of the teacher coupling for LMD and EMD.
C.4 Relation to progressive distillation
- Progressive distillation trains a new model to approximate two DDIM sampler steps with one, then repeats the process to halve the step count successively.The time increment increases across iterations until a few-step model remains.Source: §C.4 · p. 30
- LMD performs better on KL distributional correctness, whereas EMD better preserves the teacher coupling.The comparison also reports that LMD performs worse on expected L2 reconstruction error; Table 3 compares KL(ρ||ρ̂) and W2.Source: §C.4 · p. 31, Table 3 · p. 31
D.1 2D checkerboard
On the 2D checkerboard, LMD is substantially better than its Eulerian counterpart at minimizing target-to-model KL divergence. The methods trade distributional fit against preservation of the teacher’s coupling.
D.1 2D checkerboard
- LMD substantially outperforms EMD in minimizing KL divergence between the target checkerboard density and the model density.The comparison is reported with Figure 3 and Table 3.Source: §D.1 · p. 31
- LMD is more correct in distribution, while EMD better preserves the teacher’s original coupling, as measured by W2 distance and expected L2 reconstruction error.The coupling difference is visualized in Figure 6.Source: §D.1 · p. 31
- LMD more often sends initial conditions to the checker edge opposite the teacher’s destination; this is less common with EMD, whose overall target-distribution fit is worse.The figure colors points whose squared map difference from the teacher exceeds 1.0.Source: §D.1 · p. 31
D.2 Image experiments
The image experiments use a shared U-Net setup for CIFAR-10 and ImageNet-32 and visualize how model flow maps assign outputs relative to a stochastic-interpolant teacher.
D.2 Image experiments
- Both datasets use the same U-Net for velocity-field and flow-map training, with the time embedding expanded to handle two times, (s, t), for flow maps.The training conditions and image-dataset hyperparameters and architecture are presented in Table 4.Source: §D.2 · p. 31, Table 4 · p. 32
- The figure compares map assignments for the same initial conditions, marking points whose assignments differ from the teacher by more than 1.0.Blue points show where the teacher maps the red initial conditions.Source: Figure 6 · p. 32, §D.2 · p. 32
- FMM models may be trained without a teacher or with an FMM model at N = 4 as teacher, but are expected to share the same coupling at the minimizer.Source: Figure 6 · p. 32