Inductive Moment Matching

Linqi Zhou; Stefano Ermon; Jiaming Song

Source-linked AI summary

Inductive Moment Matching

Linqi Zhou, Stefano Ermon, Jiaming Song

arXiv:2503.07565v7Mar 10, 2025cs.LGcs.AIstat.ML

TL;DR

Diffusion models and Flow Matching offer high-quality samples but require slow inference, while few-step distillation can be unstable and tuning-intensive. IMM addresses this with single-stage, from-scratch training for one- or few-step generation. It guarantees distribution-level convergence and reports strong ImageNet and CIFAR-10 results.

Problem
High-quality generative modeling remains difficult to reconcile with efficient inference and stable training, while existing step-reduction approaches can require extensive tuning.
Method
IMM uses a stable, single-stage procedure with marginal-preserving interpolants and inductive distribution matching to learn one- or few-step models from scratch.
Results
IMM surpasses diffusion models on ImageNet-256×256 with 1.99 FID in 8 steps and achieves 1.98 FID on CIFAR-10 in 2 steps from scratch.
Takeaways & Limitations
IMM provides a from-scratch framework for few-step generation with distribution-level convergence and fast inference across the reported benchmarks.
Takeaways & Limitations
The single-particle estimate omits the repulsion force, while the energy kernel matches only the first moment and ignores higher moments.

Abstract

from arXiv · show

Diffusion models and Flow Matching generate high-quality samples but are slow at inference, and distilling them into few-step models often leads to instability and extensive tuning. To resolve these trade-offs, we propose Inductive Moment Matching (IMM), a new class of generative models for one- or few-step sampling with a single-stage training procedure. Unlike distillation, IMM does not require pre-training initialization and optimization of two networks; and unlike Consistency Models, IMM guarantees distribution-level convergence and remains stable under various hyperparameters and standard model architectures. IMM surpasses diffusion models on ImageNet-256x256 with 1.99 FID using only 8 inference steps and achieves state-of-the-art 2-step FID of 1.98 on CIFAR-10 for a model trained from scratch.

1. Introduction

IMM addresses the trade-off between high-quality generation, efficient inference, and stable training with a stable, single-stage procedure that learns from scratch. It supports one- or few-step sampling and achieves strong benchmark results.

IMM is a stable, single-stage training procedure that learns generative models from scratch for single- or multi-step inference.
IMM maps any time-dependent marginal at t to a marginal at s < t, enabling direct one-step and recursive few-step generation.
Inductive training matches distributions produced from different starting times, enforcing time-independent marginals and guaranteeing convergence to the data distribution.
1.99 FID on ImageNet-256×256 is achieved with 8 inference steps using standard transformer architectures, surpassing diffusion models.
1.98 FID on CIFAR-10 is achieved with 2-step generation by a model trained from scratch, reaching state-of-the-art performance.

2. Preliminaries

Diffusion models and Flow Matching construct time-augmented variables by interpolating data and prior distributions. Stochastic interpolants unify these approaches, while MMD supplies a distributional comparison based on kernel embeddings.

Diffusion and Flow Matching define xt as an interpolation between data x and Gaussian prior ϵ, with different choices of αt and σt.
Diffusion and Flow Matching train neural networks to approximate conditional interpolant velocities, then generate samples through a probability-flow ODE.
Stochastic interpolants generalize both methods by connecting arbitrary data and prior distributions through conditional Gaussian interpolations.
When the interpolant noise vanishes, the framework reduces to Flow Matching; with Gaussian prior noise, it reduces to v-prediction diffusion.
Figure 2 depicts direct transitions from qt(xt) to qs(xs), optionally composed through an intermediate r for multi-step sampling.
MMD is an RKHS-based integral probability metric that compares distributions through a positive-definite kernel and can encode moment differences.

3. Inductive Moment Matching

IMM learns an implicit sampler that transforms interpolated marginals while preserving their distributions across time. Inductive bootstrapping trains this sampler by matching recursively generated distributions, with convergence established under stated assumptions.

3.1. Model Construction via Interpolants: IMM learns a one-step sampler that transforms qt(xt) into qs(xs) for any s ≤ t using marginal-preserving interpolants.
3.1. Model Construction via Interpolants: Marginal-preserving interpolants ensure the distribution at s is independent of the starting time t.
3.1. Model Construction via Interpolants: Recursive application enables multi-step sampling by first generating an intermediate sample and then mapping it to the clean distribution.
3.2. Learning via Inductive Bootstrapping: The naïve divergence objective is difficult to optimize when qt and qs are far apart, motivating inductive bootstrapping through nearby time steps.
3.2. Learning via Inductive Bootstrapping: IMM minimizes MMD between distributions generated from t and an intermediate r, using a weighting function and an inductive sequence of objectives.
3.2. Learning via Inductive Bootstrapping: Under well-conditioned mappings, marginal-preserving interpolants, and ideal optimization assumptions, the inductive objective converges to the target distribution.

4. Simplified Formulation and Practice

IMM’s practical formulation uses self-consistent interpolants, deterministic DDIM-based mappings, particle-based MMD estimation, and several sampling options. The implementation choices cover training organization, network parameterization, kernels, weighting, and inference procedures.

4.1. Algorithmic Considerations: Self-consistent interpolants make direct and intermediate sampling produce the same xs distribution, implying marginal preservation.
4.1. Algorithmic Considerations: DDIM provides a self-consistent deterministic interpolant, allowing a deterministic minimizer that can attain zero loss under mild assumptions.
4.1. Algorithmic Considerations: The deterministic model uses DDIM to map xt to xs through gθ, and reusing xt for xr reduces variance while preserving qr(xr).
4.1. Algorithmic Considerations: IMM estimates MMD with M particle samples grouped by shared (s, t), while keeping the number of model passes independent of M.
4.2. Other Implementation Choices: Laplace kernels, time-dependent weighting, uniform time sampling, and EDM-style network coefficients are among the reported implementation choices.
4.3. Sampling: Pushforward sampling applies the learned mapping directly and can be iterated for arbitrary step counts; it is equivalent to DDIM using gθ as the x prediction.
4.3. Sampling: Restart sampling adds stochasticity by re-noising to a higher noise level, while classifier-free guidance reweights conditional outputs during inference.
4.3. Sampling: Training samples data, labels, prior noise, and shared time triples, then optimizes the empirical IMM loss with grouped particles and label dropout.

5. Connection with Prior Works

IMM reframes consistency objectives as special cases of its moment-matching framework while explaining their limitations and connecting its kernels to broader distribution-matching methods.

Consistency Models: CM objectives with L2 distance are single-particle estimates of IMM’s objective with an energy kernel.This establishes a formal connection between consistency training and IMM.
Consistency Models: Single-particle energy-kernel estimates omit entropy repulsion and match only the first moment, contributing to CM instability and mode collapse.The energy kernel is not a proper positive definite MMD kernel, unlike kernels that capture higher moments.
Kernel Design: Pseudo-Huber loss is a conditionally positive definite kernel whose higher-moment weighting depends on c.This provides a valid alternative to the energy kernel within the moment-matching design space.
Consistency Models: IMM’s differential-limit loss subsumes continuous-time Consistency Models as a single-particle estimate, although experiments are left for future work.The connection is stated theoretically rather than empirically evaluated here.
Adversarial Methods: Diffusion GAN and Adversarial Consistency Distillation use GAN objectives that share similarity with MMD.IMM’s moment-matching perspective places these adversarial approaches within a related distribution-matching landscape.
Generative Moment Matching Network: GMMN is a special case of IMM when t = 1 and r(s, t) ≡ s = 0, reducing IMM’s loss to the naïve GMMN objective.Both methods use moment-based distribution matching, but IMM generalizes the setting to intermediate marginals.

6. Related Works

Prior work includes diffusion, Flow Matching, stochastic interpolants, diffusion distillation, and few-step models, with IMM positioned as a direct mapping approach between intermediate marginals.

Diffusion, Flow Matching, and stochastic interpolants: Diffusion and Flow Matching learn score or velocity fields that transfer data distributions toward a simple prior.These frameworks have been successfully scaled to text-to-image generation but remain associated with sampling inefficiency in the surrounding motivation.
Diffusion, Flow Matching, and stochastic interpolants: Stochastic interpolants explicitly define paths between data and prior distributions and match their velocity to facilitate distribution transfer.IMM builds on this construction while directly learning one-step mappings between intermediate marginal distributions.
Diffusion distillation: Diffusion distillation trains one- or few-step models from pretrained diffusion models, but jointly optimized two-network approaches require careful tuning and can lead to mode collapse.This motivates alternatives that avoid distillation’s training requirements.
Few-step generative models from scratch: Early few-step generative models relied mainly on GANs, MMD, or combinations of both, while scaling adversarial training remained challenging.Recent independent classes include Consistency Models, Consistency Trajectory Models, and Shortcut Models.

7. Experiments

Experiments evaluate IMM’s image quality, stability, sampling, scaling, and parameter choices. IMM achieves strong few-step FID while remaining stable across several practical settings.

7.1. Image Generation: 1.99 FID is achieved on ImageNet-256×256 with 8 inference steps, surpassing DiT and SiT with the same architecture.
7.2. IMM Training is Stable: Both Fourier and positional embeddings converge reliably, unlike the reported Fourier-embedding instability in Consistency Models.
7.2. IMM Training is Stable: M = 1 and M = 2 collapse on ImageNet-256×256, while M = 4 achieves the lowest FID under the same computation budget.
7.3. Sampling: Uniform pushforward schedules work best among the investigated ImageNet-256×256 settings, although restart samplers can achieve superior results.
7.4. Scaling Behavior: IMM’s FID and sample quality improve with training compute, inference compute, model size, and inference steps.
7.5. Ablation Studies: On larger ImageNet-256×256 models, Euler parameterization with OT-FM performs best, while identity parameterization consistently lags.
7.5. Ablation Studies: Constant decrement in ηt consistently outperforms the other tested mapping functions, while k ∈ {11, 12, 13} is stable and k = 14 is unstable.

8. Conclusion

The conclusion presents Inductive Moment Matching as a from-scratch framework for few-step generation. It combines self-consistent interpolation with all-moment matching, guarantees distributional convergence, and reports state-of-the-art benchmark performance with much faster inference.

IMM learns a few-step generative model from scratch by interpolating between data and prior and matching all moments toward the data distribution.
IMM guarantees convergence in distribution, generalizes many prior works, and achieves state-of-the-art performance across benchmarks with orders-of-magnitude faster inference.

Impact Statement

The paper situates generative modeling within unresolved technical challenges and discusses broader societal benefits and risks of advancing diffusion-based systems.

The paper identifies potential benefits including artistic expression, content-creator assistance, and synthetic data generation for research.
The paper acknowledges potential misuse for deepfakes, copyright concerns, and impacts on creative industries.
The discussion encourages continued consideration of responsible development and deployment alongside technical progress.

B.4. Main Theorem

The main theorem establishes inductive convergence under well-conditioned mappings, marginal-preserving interpolants, and ideal optimization assumptions, while self-consistency preserves intermediate marginals.

Under the theorem’s assumptions, each inductive parameter θ_n minimizes the corresponding objective over all valid times s and t.The assumptions include a well-conditioned r(s,t), a marginal-preserving interpolant, infinite data and network capacity, and existence of θ* minimizers.
The induction repeatedly applies the inverse time mapping and covers the entire interval from s to 1.
As n approaches infinity, the model distribution converges to the target marginal in MMD.
Self-consistent interpolants make the induced marginal at s invariant to the starting time t.
Deterministic minimizers exist under invertibility, zero conditional noise, and a bounded regularity condition, including DDIM-based interpolants.

C. Analysis of Simplified Parameterization

The simplified parameterization analyzes self-consistent DDIM transitions, variance reduction through sample reuse, and practical kernel and weighting choices for stable IMM optimization.

DDIM interpolant: DDIM is self-consistent because composing transitions from t to r and r to s equals the direct transition from t to s.
Variance reduction: Reusing x_t to construct x_r preserves the marginal distribution while reducing sampling variance.
MMD estimation: IMM’s particle-based MMD estimate groups samples sharing (s, r, t) and requires only two model forward passes regardless of particle count M.
Parameterizations: The parameterizations connect x-prediction, v-prediction, and epsilon-prediction differently under cosine and OT-FM schedules.
Time mappings: A constant decrement in η(t) is preferred empirically, with k between 10 and 15 working across datasets.
Kernel choice: The Laplace kernel self-normalizes its gradient, unlike the RBF gradient whose magnitude can vary with point distance.
Loss weighting: The α_t^2 weighting implicitly produces an epsilon-prediction parameterization and emphasizes smaller time-steps more than α_t weighting.

G.1. Consistency Models

The analysis relates IMM to consistency models, diffusion GANs, GMMN, and continuous-time consistency objectives, while identifying divergence and kernel distinctions.

Consistency Models: Consistency-model loss appears as a single-particle special case of IMM under a small positive s and a negative L2 kernel.
Consistency Models: The consistency-model loss is not a proper distribution divergence because single-particle estimation omits MMD’s entropy repulsion term.
Consistency Models: The energy kernel used there is not positive definite and, at best, matches only the first moment.
Kernel analysis: The pseudo-Huber kernel is conditionally positive definite and its expansion contains terms corresponding to all higher moments.
Diffusion GAN: DDPM posterior distributions are self-consistent Gaussian interpolants, providing another interpolant design compatible with the framework.
Related methods: MMD avoids adversarial optimization by selecting an RKHS whose optimal function can be found analytically, yielding a stable minimization objective in practice.
GMMN: GMMN is a special IMM parameterization obtained by fixing t = 1 and mapping s to 0.

I.1. Training & Parameterization Settings

The experiments use latent-space ImageNet models with standardized architectural and schedule settings, while evaluating inference schedules, compute, scaling, and lower-precision adaptations.

Experimental setup: ImageNet experiments use pre-encoded, normalized Stable Diffusion VAE2 latents and shared hyperparameters across DiT model sizes.Latents are rescaled by channel statistics and multiplied by 0.5 to reach roughly standard deviation 0.5.
Inference schedules: One-step sampling uses pushforward from T to ϵ, while two-step sampling sets ηt1 = 1.4 empirically.For at least four steps, the experiments compare uniform t decrements with EDM schedules.
Scaling beyond 8 steps: 16 steps achieve 1.92 FID on ImageNet-256×256, already outperforming VAR with 2B parameters, while performance saturates beyond that.The reported 32-step result is only marginally better.
Parameterization and precision: The weighting exponent a = 2 slightly worsens one-step FID but outperforms a = 1 in the multi-step regime.For lower-precision training, a minimum gap Δ = 10^-4 and (t − r) conditioning address indistinguishable time embeddings.