PaperlayerView exact revision on arXiv

Source-linked AI summary

How to build a consistency model: Learning flow maps via self-distillation

Nicholas M. Boffi, Michael S. Albergo, Eric Vanden-Eijnden

arXiv:2505.18825v2cs.LGcs.CV

TL;DR

Flow-based generative models require expensive differential-equation solves at inference, while existing flow-map methods lack a unified efficient training methodology. The paper introduces self-distillation for direct flow-map training through three mathematical characterizations and finds that Lagrangian methods provide more stable, higher-performing training. Flow maps thereby offer a unified approach to accelerated inference without pre-trained teachers.

  • Problem

    Flow-based generative models require costly differential-equation solutions at inference, and flow-map methods lack a unified methodology for efficient training.

  • Method

    The paper develops a mathematical self-distillation framework that directly trains flow maps using the tangent condition and three equivalent characterizations, encompassing existing methods.

  • Results

    LSD attains the best performance across standard benchmarks, while theoretical guarantees bound the 2-Wasserstein error for ESD and LSD.

  • Takeaways & Limitations

    Flow maps can be trained directly with minimal additional overhead, avoiding differential-equation solves during inference and enabling accelerated generative modeling.

  • Takeaways & Limitations

    The paper focuses on one parameterization and leaves more general and higher-order parameterizations to future work.

Abstract

from arXiv · show

Flow-based generative models achieve state-of-the-art sample quality, but require the expensive solution of a differential equation at inference time. Flow map models, commonly known as consistency models, encompass many recent efforts to improve inference-time efficiency by learning the solution operator of this differential equation. Yet despite their promise, these models lack a unified description that clearly explains how to learn them efficiently in practice. Here, building on the methodology proposed in Boffi et. al. (2024), we present a systematic algorithmic framework for directly learning the flow map associated with a flow or diffusion model. By exploiting a relationship between the velocity field underlying a continuous-time flow and the instantaneous rate of change of the flow map, we show how to convert any distillation scheme into a direct training algorithm via self-distillation, eliminating the need for pre-trained teachers. We introduce three algorithmic families based on different mathematical characterizations of the flow map: Eulerian, Lagrangian, and Progressive methods, which we show encompass and extend all known distillation and direct training schemes for consistency models. We find that the novel class of Lagrangian methods, which avoid both spatial derivatives and bootstrapping from small steps by design, achieve significantly more stable training and higher performance than more standard Eulerian and Progressive schemes. Our methodology unifies existing training schemes under a single common framework and reveals new design principles for accelerated generative modeling. Associated code is available at https://github.com/nmboffi/flow-maps.

1 Introduction

Flow-based generative models offer strong sample quality but require costly differential-equation solutions at inference. This work develops a unified framework for directly training flow maps through self-distillation and introduces three algorithmic families, with LSD achieving the best benchmark performance.

  • Differential-equation solution during generation creates a computational bottleneck for flows and diffusions, limiting applications requiring rapid inference.
  • Distillation methods perform well but require pre-trained teachers, two-phase training, and remain limited by teacher performance.
  • The proposed tangent condition relates probability-flow velocity to flow-map derivatives, enabling direct self-distillation without pre-trained teacher models.
  • The framework gives three equivalent flow-map characterizations and shows that consistency models and several recent few-step methods are special cases.
  • The new Eulerian, Lagrangian, and Progressive Self-Distillation families provide distinct algorithms for learning flow maps.
  • The framework's theoretical guarantees show that loss values bound the 2-Wasserstein error of learned one-step models for ESD and LSD.
  • LSD attains the best performance across standard benchmarks, including checkerboard, CIFAR-10, CelebA-64, and AFHQ-64.

2 Theoretical framework

The paper characterizes flow maps as operators that jump along probability-flow trajectories and derives self-distillation objectives from their mathematical properties. This framework supports direct training without pre-trained teachers and provides accuracy guarantees for Eulerian and Lagrangian methods.

  • Flow-map definition: Flow maps jump arbitrary distances along probability-flow trajectories, allowing one-step generation without numerical integration or multistep sampling when model errors are present.A single X0,1 maps base samples to the target; composing smaller jumps trades inference-time computation for sample quality.
  • Flow-map characterization: The tangent condition identifies the velocity field as the instantaneous time derivative of the flow map, linking flow-map learning to flow matching.The Euler step-like parameterization enforces Xs,s(x) = x and connects vt,t to bt.
  • Flow-map characterization: Eulerian, Lagrangian, and semigroup conditions provide three equivalent characterizations of the flow map and motivate distinct estimation objectives.The conditions describe transport through PDE, trajectory-wise ODE, and compositional formulations.
  • Self-distillation: Self-distillation combines diagonal flow-matching training with off-diagonal penalties enforcing flow-map conditions, yielding direct training objectives for LSD, ESD, and PSD.The parameterization Xs,t(x) = x + (t −s)vs,t(x) uses a single learned function and supports these objectives.
  • Guarantees: Each ESD and LSD objective has a unique correct minimizer, and its loss value bounds the 2-Wasserstein error of the learned one-step model.The guarantees connect optimization of the combined flow-matching and distillation losses to model accuracy.

3 Algorithmic aspects

The algorithmic framework offers practical training choices for self-distillation, including loss weighting, temporal sampling, parameterization, and teacher design. It focuses on stable direct training while identifying capacity and bootstrapping considerations.

  • Connections to existing methods: The framework recovers standard consistency, consistency trajectory, and shortcut models through suitable choices of distillation objective and teacher.These existing methods appear as special cases of the generic self-distillation framework.
  • Loss weighting: Learned two-time loss weights normalize variation across time pairs, stabilizing training dynamics and enabling larger learning rates.The weights estimate loss log-variance and make different (s, t) contributions comparable near the global minimizer.
  • Temporal sampling: Temporal sampling mixes diagonal flow-matching examples with off-diagonal distillation examples to train both components of the self-distillation objective.The mixture uses Ud on s = t and Uod on s < t, with fraction η assigned to the diagonal.
  • PSD scaling: Dropping the parameterization-induced (t −s)2 factor preconditions the PSD loss and improves training stability by removing timestep-dependent gradient variance.Without this factor, the effective learning rate no longer changes solely because of the time gap.
  • Model parameterization: A single model jointly representing velocity and flow map is efficient, but may require more network capacity than using separate models.The paper also describes two-model and higher-order alternatives, leaving them for future work.

4 Related work

The paper situates flow maps within flow matching, diffusion, and consistency-model research aimed at reducing the cost of differential-equation-based sampling. It distinguishes single-time and two-time flow-map approaches and reports benchmark evaluation conventions.

  • Flow matching and diffusion: Flow matching, stochastic interpolants, and diffusion probability-flow equations define dynamics whose numerical solution evaluates the flow map but requires repeated model evaluations.The paper builds its framework directly on these dynamical generative-model formulations.
  • Benchmark evaluation: Table 1 compares KL divergence on the checker dataset and FID on natural-image datasets across sampling step counts, with the best method marked per dataset and count.The caption defines the metrics, datasets, and comparison layout.
  • Consistency models: Consistency models estimate a single-time map Xs,1 for direct jumps to data, whereas consistency trajectory models estimate two-time maps that enable multistep sampling.These methods are presented as related approaches to accelerated inference.

5 Numerical experiments

Experiments compare LSD, ESD, and PSD variants on checkerboard and image-generation tasks. LSD achieves the strongest overall performance and more stable training, while increasing sampling steps generally improves quality.

  • Experimental setup and overall findings: LSD obtains the best performance across the checkerboard, CIFAR-10, CelebA-64, and AFHQ-64 experiments, with the checkerboard exception at N = 16.The study evaluates methods at fixed training time for equal-footing comparisons.
  • Checkerboard: On the checkerboard, ESD and PSD variants miss sharp boundaries at small N, introducing artifacts and increasing KL divergence.All methods improve as the number of steps increases, and LSD performs best except at N = 16.
  • CIFAR-10: ESD becomes unstable on CIFAR-10 because its high parameter gradient norms induce training divergence.LSD and PSD avoid spatial derivatives and maintain more stable gradients, even with stopgradients.
  • Image datasets: LSD achieves the best FID across all reported CIFAR-10, CelebA-64, and AFHQ-64 step counts, while ESD has no reported CIFAR-10 FID because of instability.On CelebA-64, LSD ranges from FID 12.22 at N = 1 to 1.96 at N = 16; on AFHQ-64, from 11.19 to 5.61.
  • Progressive refinement: Increasing sampling steps improves generated-image quality, producing sharper and more detailed CelebA-64 faces and increasingly detailed AFHQ-64 images.CIFAR-10 methods also improve with more steps, while fixed-seed samples remain similar across methods.

6 Conclusion

The paper develops a mathematical and algorithmic framework for learning flow maps through self-distillation and connects probability-flow dynamics to efficient generative sampling. It reviews the transport formulation and teacher-based distillation, then replaces the frozen teacher with an implicit self-consistent teacher.

  • Probability-flow background: The interpolant density satisfies a transport equation whose drift is the conditional expectation of the interpolant velocity.The corresponding probability-flow ODE has the same time-marginal law as the interpolant.
  • Self-distillation: Existing distillation schemes use a frozen pre-trained velocity teacher, whereas the proposed self-distillation schemes use a self-consistent implicit teacher.Stopgradients emulate a frozen teacher and control information flow during training.
  • Flow-map characterizations: The framework derives training schemes from three equivalent flow-map characterizations: Lagrangian, Eulerian, and semigroup conditions.These characterize transport along trajectories, a PDE description, and replacement of successive jumps by one jump.

B.1 Lagrangian distillation.

Lagrangian map distillation trains a flow map by matching its time derivative to a velocity evaluated at a transported point. The approach motivates self-distillation variants that avoid dependence on a pre-trained teacher.

  • Lagrangian map distillation: Lagrangian map distillation matches the flow-map time derivative with a velocity evaluated at the transported point.The resulting loss is nonconvex in the flow map and requires the velocity field’s spatial Jacobian during backpropagation.
  • Lagrangian map distillation: Spatial-Jacobian backpropagation can be problematic for large generative models, motivating a semigradient modification with stopgrad.The effectiveness of this modification depends on the data modality and neural-network architecture.
  • Eulerian map distillation: Eulerian map distillation instead uses an objective convex in the flow map, but its parameter gradients still require backpropagation through a spatial Jacobian.A semigradient variant avoids that backpropagation entirely.
  • Comparison: Lagrangian schemes were observed to be more stable than Eulerian schemes, consistent with the experiments reported in the paper.The comparison concerns the Lagrangian losses (26) and (27) versus the Eulerian losses (28) and (29).
  • Consistency connections: Continuous-time consistency distillation emerges from the Eulerian objective, while consistency training replaces teacher trajectory points with two points from the same interpolant trajectory.The resulting semigradient is equivalent to the continuous-time consistency-distillation objective after rescaling gradients.

C.2 Consistency training and mean flow.

Consistency training removes the need for a pre-trained teacher by using nearby points on the same interpolant trajectory. Its mean-flow formulation trades lower bias for higher-variance Monte Carlo estimates.

  • Consistency training: Consistency training avoids a pre-trained teacher by comparing two points on the same interpolant trajectory.This replaces a nearby point generated using the teacher velocity in consistency distillation.
  • Consistency training: The resulting Monte Carlo update substitutes the interpolant velocity for the true vector field in the consistency objective.Shared endpoints between the two interpolant points produce the required local expansion.
  • Mean flow: Mean-flow and related consistency-training semigradients have higher Monte Carlo variance but access the ideal flow rather than a pre-trained approximate flow.The mean-flow update is identified with the gradient in equation (42).
  • Shortcut models: Shortcut models are represented as a discretized progressive self-distillation scheme using a fixed dyadic time grid and bootstrapping.Their objective matches the progressive formulation under restricted times and a specific proposal distribution.

E Proofs

The proofs establish that Lagrangian, Eulerian, and semigroup characterizations identify the same flow map under stated regularity assumptions. They then show that the self-distillation objectives have the ideal map as their unique minimizer, with error bounds for LSD and ESD.

  • Assumptions: Under a one-sided Lipschitz assumption, the probability-flow differential equation has unique solutions for all initial points and times.The assumption is stated as (44), and classical Cauchy-Lipschitz theory supplies existence and uniqueness.
  • Tangent condition: The tangent condition identifies the velocity field with the instantaneous time derivative of the flow map along its trajectory.This relation is the mathematical basis for the self-distillation construction.
  • Flow-map characterizations: The flow map satisfies Lagrangian, Eulerian, and semigroup conditions, each providing an equivalent characterization under the paper’s assumptions.The semigroup property composes two jumps into one.
  • Self-distillation guarantees: Each LSD, ESD, and PSD objective has the ideal flow map as its unique global minimizer under the stated continuity and regularity conditions.The proofs establish optimality through the corresponding Lagrangian, Eulerian, or semigroup condition.

F.6 Detailed algorithms for each self-distillation method

The detailed algorithms combine diagonal flow matching with off-diagonal self-distillation and update a shared velocity model whose parameterization yields the trained flow map. LSD, ESD, and PSD differ in the teacher signals and intermediate computations they use.

  • Shared structure: All three algorithms combine diagonal flow matching with a distillation loss and output a flow map parameterized as ˆX_s,t(x) = x + (t −s)ˆv_s,t(x).The implementation uses separate diagonal and distillation batches within a total batch size M.
  • Lagrangian Self-Distillation (LSD): LSD samples diagonal times and off-diagonal time pairs, computes a flow-map derivative and transported teacher velocity, then updates the shared model.Its output is ˆX_s,t(x) = x + (t −s)ˆv_s,t(x).
  • Eulerian Self-Distillation (ESD): ESD uses the same diagonal and off-diagonal batching structure while computing the Eulerian self-distillation loss.The algorithm is specified as Eulerian Self-Distillation with the shared model and flow-map output.
  • Progressive Self-Distillation (PSD): PSD samples intermediate times, composes two predicted flow-map segments, and uses the resulting second-segment velocity as a preconditioned teacher signal.The model is trained with the flow-matching loss plus the progressive self-distillation loss.

G Further details on numerical experiments

This section provides a complete description of the numerical experiments and summarizes each experiment in Table 2.

  • The section gives a complete description of the numerical experiments.
  • Each experiment receives a concise summary in Table 2.
  • Table 2 serves as the centralized summary of the experiments described in this section.

G.1 Checkerboard Details

The checkerboard experiments compare the three training algorithms under specified sampling, optimization, and evaluation procedures. KL divergence is estimated using the analytically known target density, histograms, and deterministic quadrature.

  • Experimental setup: The checkerboard comparison evaluates LSD, ESD, PSD-U, and PSD-M using a 10^7-sample dataset.PSD-U uses uniform γ sampling, whereas PSD-M uses midpoint sampling.
  • Experimental setup: Training runs for 150,000 steps with batch size 100,000, learning rate 10−3, and square-root decay after 35,000 steps.Each batch allocates 75% to flow matching and 25% to self-distillation.
  • KL computation: KL divergence uses the analytically known uniform checkerboard density as the target distribution.This provides the reference density for evaluating model samples.
  • KL computation: The model density is approximated by histogramming samples because the checkerboard dataset is low-dimensional.
  • KL computation: The KL estimate uses quadrature over a 50 × 50 grid spanning [−1, 1]2.The grid is applied after computing 64,000 samples from each model for each number of steps.
  • KL computation: Because both densities are uniform in the relevant construction, the quadrature rule is exact.

G.2 CIFAR-10 Details.

The supplied experiments evaluate LSD, ESD, and PSD variants on CIFAR-10, CelebA-64, and AFHQ-64 using dataset-specific training configurations and sample-quality evaluation procedures.

  • CIFAR-10: CIFAR-10 experiments compare LSD, ESD, PSD-U, and PSD-M with uniform upper-triangle sampling and no annealing.Training uses 400,000 steps, batch size 512, and an initial learning rate of 10−2.
  • CIFAR-10: CIFAR-10 uses a learned 128-channel weight function, adaptive Gaussian variance, gradient clipping at 1.0, and specified stopgradient configurations.ESD was unstable in every tested stopgradient configuration.
  • CelebA-64: CelebA-64 compares LSD with uniform- and midpoint-sampling PSD variants, while ESD is omitted because it was uniformly unstable.
  • CelebA-64: CelebA-64 training runs for 800,000 steps with batch size 256 and an initial learning rate of 10−2.The setup allocates 75% of each batch to flow matching and 25% to self-distillation.
  • CelebA-64: CelebA-64 evaluates online FID-10K every 10,000 steps and computes FID-50k post hoc for the best model.Models are trained from random initialization without pre-training.
  • AFHQ-64: AFHQ-64 compares LSD with both PSD variants under an 800,000-step training configuration.As with CelebA-64, ESD was unstable and is not reported.
Loading 2505.18825v2…