Inductive Moment Matching

Inductive Moment Matching

Speakers: Linqi (Alex) Zhou; Stefano Ermon; Jiaming Song

Conference: ICML 2025

YouTube: https://slideslive.com/39044001

Overview

In the rapidly evolving landscape of generative AI, particularly in visual domains, diffusion models and flow matching have emerged as dominant paradigms, powering sophisticated text-to-image and text-to-video systems. While these models have achieved unprecedented levels of realism, they frequently grapple with a fundamental challenge often termed the "generative trilemma": the difficulty in simultaneously achieving high sample quality, stable training, and efficient inference. This talk introduces Inductive Moment Matching (IMM), a novel approach designed to address this trilemma head-on, offering a single-stage training objective that promises to deliver all three desirable properties.

Presented by Linqi (Alex) Zhou alongside his advisor Stefano Ermon and Luma AI's Chief Scientist Jiaming Song, this work presents IMM as a significant step towards a new class of generative models. The core innovation lies in enhancing the capacity of inference-time samplers like DDIM and introducing a stable, inductive training mechanism based on Maximum Mean Discrepancy (MMD). The implications for practitioners are substantial, offering a pathway to build and deploy high-quality generative models with faster iteration cycles and reduced computational overhead, without compromising on the fidelity of generated outputs.

The significance of IMM stems from its ambition to overcome long-standing tradeoffs that have plagued generative model development. By demonstrating superior performance in terms of sample quality, training stability, and inference efficiency compared to established baselines, IMM positions itself as a robust alternative. Its ability to scale effectively with computational resources and model size further underscores its potential to influence the next generation of large-scale foundation models, particularly in domains like realistic video generation, where efficiency and quality are paramount.

Background

▶ Watch: Introduction and the generative model trilemma (0:00)

The past half-decade has seen an explosion in the capabilities of generative models, with diffusion models and flow matching techniques leading the charge in continuous domains. Architectures like Stable Diffusion and Midjourney have revolutionized creative industries by generating hyper-realistic images from text prompts, while emerging text-to-video models push the boundaries of dynamic content creation. Despite these successes, the underlying methodologies often present a challenging compromise, commonly referred to as the "generative trilemma."

This trilemma highlights that no single existing generative model class perfectly achieves all three desirable properties:

1. High Sample Quality: The ability to produce outputs that are indistinguishable from real data.
2. Stable Training: A robust and predictable optimization process that reliably converges.
3. Efficient Inference: The capacity to generate samples quickly, ideally in a single or very few steps.

For instance, diffusion models excel at producing high-quality samples and boast stable training, but their inference process is notoriously slow, requiring hundreds or even thousands of sequential steps to generate a single sample. Generative Adversarial Networks (GANs) and diffusion distillation methods, while fast at inference and capable of high quality, frequently suffer from training instability, making them challenging to optimize. Variational Autoencoders (VAEs) and Normalizing Flows, on the other hand, offer stable training but often lag in sample quality and can still be slow for complex generation tasks. This fragmentation means that developers often have to choose which properties to prioritize, leading to suboptimal solutions for real-world applications. The motivation behind IMM is to transcend these limitations, offering a unified framework that simultaneously addresses all three facets of the generative trilemma under a single training stage and objective.

Key Findings

▶ Watch: Introducing Inductive Moment Matching (IMM) components (4:00)

Inductive Moment Matching (IMM) introduces a principled approach to resolve the generative trilemma, demonstrating significant advancements across all three critical axes: inference efficiency, training stability, and sample quality. The core findings highlight IMM's ability to:

1. Achieve Inference Efficiency through Enhanced Capacity: IMM fundamentally re-architects the inference process by injecting the target time step S directly into the neural network, alongside the current state XT and current time T. This allows the network, leveraging its universal function approximation capabilities, to learn complex, non-linear mappings that can perform large, multi-step jumps in the generative process. This contrasts sharply with traditional DDIM samplers, which are limited to linear relationships and infinitesimal steps, thereby enabling IMM to achieve high-quality samples in as few as 8 or 16 inference steps.

1. Ensure Stable Training via Inductive Maximum Mean Discrepancy (MMD): IMM proposes a novel training objective comprising two key components. First, it employs a sample-based distribution matching objective based on Maximum Mean Discrepancy (MMD). MMD offers a stable alternative to adversarial training by implicitly matching all moments of two distributions within a Reproducing Kernel Hilbert Space (RKHS), thus avoiding the instability issues common in GANs. Second, IMM incorporates inductive learning, where the model's own distribution at an intermediate, close-by state (QR to PSR) serves as the training target, rather than directly matching a potentially distant ground truth distribution (QS). This inductive assumption simplifies the learning task, making the optimization process inherently more stable and requiring fewer samples for accurate moment estimation.

1. Produce High-Quality Samples with Competitive FID Scores: Empirically, IMM demonstrates superior sample quality compared to established baselines. It achieves an FID (Fréchet Inception Distance) of 1.99 with only 8 inference steps and 1.90 with 16 steps when trained from scratch. These figures notably outperform diffusion and flow matching models using identical architectures. Furthermore, IMM proves competitive with, and often superior to, autoregressive baselines like VAR, achieving comparable quality with significantly fewer parameters (e.g., a 600 million parameter IMM model matching a 2 billion parameter VAR model) and offering flexible control over inference steps.

1. Scalability and Flexibility: The method scales effectively with both training compute and sampling compute. Visual results consistently show improved sample quality with an increased number of sampling steps and larger transformer sizes, indicating robust scalability. The flexibility to choose the number of inference steps at generation time is a distinct advantage, allowing users to balance speed and quality according to their specific needs.

Technical Deep Dive

▶ Watch: Sample-based distribution matching using MMD (4:25)

Inductive Moment Matching (IMM) fundamentally re-evaluates the paradigms of generative model inference and training, proposing a unified framework that tackles the "generative trilemma" head-on. The technical innovations span two primary areas: enhancing inference efficiency and designing a stable, single-stage training objective.

Enhanced Inference Efficiency

The journey towards efficient inference in IMM begins with a critical assessment of existing methods, particularly those based on solving probability flow ODEs in diffusion and flow matching models. These methods, while effective, suffer from two inherent limitations: ODE simulation error and slow inference. Even with a perfectly predicted velocity field, simulating an ODE accurately necessitates numerous infinitesimal steps, which is computationally expensive and time-consuming.

The core idea behind IMM's inference efficiency is to enable the neural network to perform large, non-infinitesimal jumps across time steps, effectively learning a direct mapping from a prior distribution to the data distribution in one or very few steps. This concept is visualized as a "green line" trajectory in contrast to the many small steps of ODE simulation.

Traditional inference samplers, such as DDIM (Denoising Diffusion Implicit Models), commonly used in diffusion models, are identified as a bottleneck. A standard DDIM step, which reduces to an Euler sampler under a flow matching schedule, typically takes in the current noisy data XT and the current time T to predict the denoised data X0 or a velocity field. However, the resulting transformation XS (from QT to QS) is often limited to a linear relationship with respect to the target time S. This inherent linearity severely constrains the network's capacity to represent complex, non-linear mappings that are crucial for large time jumps.

IMM addresses this capacity issue by a simple yet powerful modification: it injects the target time step S directly into the neural network's input, alongside XT and T. So, instead of network(XT, T), the input becomes network(XT, T, S). Since neural networks are universal function approximators, providing S as an explicit input allows the network to learn a much more complex, non-linear relationship between XT, T, and S. This enhanced capacity enables the model to directly represent a one-step or few-step mapping from any intermediate time T to any other time S, effectively performing "ODE integration" in a single network call. This direct learning of large jumps dramatically speeds up sampling, resolving the inference efficiency issue.

Single-Stage Stable Training Objective: Inductive Moment Matching

The second cornerstone of IMM is its novel, single-stage training objective, which combines sample-based distribution matching with inductive learning. This design aims to ensure stable optimization while achieving high sample quality.

1. Sample-Based Distribution Matching with MMD:

IMM utilizes Maximum Mean Discrepancy (MMD) as its distribution matching objective. MMD is a powerful statistical test that measures the distance between two probability distributions based on samples from them. It offers significant advantages over adversarial training methods (like those used in GANs) because it bypasses the need for a dynamically optimized discriminator. Instead, MMD leverages a fixed, optimal discriminator implicitly chosen within a Reproducing Kernel Hilbert Space (RKHS), typically defined by a kernel function such as the Radial Basis Function (RBF) kernel or Laplace kernel.

The key benefits of MMD in this context are:

• Stability: By eliminating the adversarial min-max game, MMD-based optimization is inherently more stable and less prone to mode collapse or training oscillations.
• Moment Matching: Standard kernel functions used in MMD implicitly match all moments of the two distributions being compared, hence the name "Moment Matching." This ensures a comprehensive alignment of the learned distribution with the target.
• Empirical Estimation: The MMD objective is estimated using multiple particles (samples) to approximate the expectation, making it practical for neural network training. The talk notes that using more than four particles empirically leads to stable training.

1. Inductive Learning:

While MMD provides a stable metric, directly matching the model's distribution P(S|XT, T) to a ground truth target distribution QS can still be challenging, especially if QS is very "far away" from the current model's distribution QT. This is particularly true in extreme cases where QT is the prior and QS is the data distribution, requiring a single, enormous jump. Attempting to match moments between vastly different distributions requires a very large number of samples to accurately estimate, leading to empirical difficulties.

IMM's inductive learning component addresses this by introducing an "inductive assumption." Instead of directly matching P(S|XT, T) with QS, the training target becomes the model's own distribution at an intermediate step. Specifically, the model learns to map an intermediate marginal distribution QR (chosen to be very close to QT) to PSR (the model's output at R). The "inductive assumption" is that this PSR is already sufficiently close to the ground truth QS.

This approach simplifies the learning problem:

• Easier Task: By matching distributions that are inherently closer (e.g., QR to PSR where R is close to T), the network has an easier task. This requires fewer samples to accurately estimate the moments for the MMD objective.
• Progressive Learning: This inductive scheme facilitates a progressive learning process, where the model iteratively refines its ability to make small, accurate jumps, which then implicitly builds its capacity for larger jumps.
• Connection to Consistency Training: The speaker notes that consistency training can be viewed as a special case of IMM's inductive learning with a single particle, highlighting the broader theoretical underpinnings.

By combining the capacity-enhanced DDIM for inference with the stable, inductive MMD objective for training, IMM creates a robust and efficient generative model that simultaneously addresses the core challenges of quality, stability, and speed. The objective function is designed to minimize the MMD between the model's predicted distribution and the inductively chosen target distribution, using a single, unified training stage.

Experimental Setup & Results

▶ Watch: Inductive learning for stable training objective (5:00)

The evaluation of Inductive Moment Matching (IMM) centers on demonstrating its efficacy across the three dimensions of the generative trilemma: sample quality, inference efficiency, and training stability. The primary metric for sample quality is the Fréchet Inception Distance (FID), a widely accepted measure that quantifies the similarity between generated and real images. Lower FID scores indicate higher quality and fidelity.

Key Quantitative Results:

• Sample Quality: IMM achieves impressive FID scores with a remarkably low number of inference steps:
• 1.99 FID with only 8 inference steps.
• 1.90 FID with 16 inference steps.

These results were obtained when training IMM from scratch, using a vanilla DiT (Diffusion Transformer) architecture. This performance significantly surpasses traditional diffusion models and flow matching baselines when evaluated with the same architectural backbone and comparable computational budget.

• Comparison to Autoregressive Models: IMM also shows strong competitiveness against high-performing autoregressive models like VAR (Visual Autoregressive models). While VAR can achieve high-quality samples, IMM often matches or exceeds its quality with substantially fewer parameters. For instance, a 600 million parameter IMM model with 16 inference steps can achieve quality comparable to a 2 billion parameter VAR model. This highlights IMM's superior parameter efficiency.

• Inference Efficiency: The ability to achieve high-quality samples in 8-16 steps is a direct testament to IMM's enhanced inference efficiency. This is a dramatic improvement over diffusion models which typically require hundreds or thousands of steps for comparable quality. The flexibility to choose the number of steps at inference time further allows for a tunable trade-off between speed and quality.

• Training Stability: Empirically, IMM's training process is shown to be highly stable. The use of the MMD objective, particularly when estimated with more than four particles, avoids the common instability issues associated with adversarial training. This stability is a direct benefit of the MMD's properties as a fixed, optimal discriminator and the inductive learning approach which simplifies the distribution matching task.

Scaling and Visual Results:

• Scaling with Compute: The talk presents graphs illustrating that IMM scales effectively with both training compute and sampling compute. As the amount of computational resources (e.g., larger transformer sizes) increases, the sample quality (e.g., sharpness, detail) consistently improves.
• Step-by-Step Refinement: Visual results demonstrate that the quality of generated samples noticeably improves step-by-step with an increasing number of sampling steps, showcasing the model's ability to refine its outputs progressively. This indicates that while few steps are sufficient for good quality, more steps can push fidelity even higher.

Hardware and Datasets:

While specific dataset names and hardware configurations (e.g., GPU types, cluster sizes) are not explicitly detailed in the transcript, the context of "large-scale foundation models" and "visual generative model landscape" implies training on extensive image and video datasets (likely in the millions or billions) using state-of-the-art GPU clusters, similar to those used for diffusion models and large language models. The mention of "vanilla DiT" architecture suggests a transformer-based backbone, common in modern generative models.

In summary, IMM's experimental results underscore its ability to deliver on the promise of the generative trilemma, offering a robust, efficient, and high-quality solution that sets a new benchmark for generative modeling.

Practical Implications

▶ Watch: Training stability and connection to consistency training (6:00)

Inductive Moment Matching (IMM) carries profound practical implications for a wide spectrum of professionals involved in machine learning, from model builders and infrastructure teams to practitioners deploying AI systems. By addressing the long-standing "generative trilemma," IMM offers tangible benefits that can reshape how generative models are developed and utilized.

For Practitioners and Model Builders:

• Accelerated Development Cycles: The primary benefit is the ability to achieve high-quality generative models with significantly faster inference times. This means faster iteration during development, quicker prototyping of new ideas, and reduced time-to-market for applications. A model that can generate high-fidelity samples in 8-16 steps, rather than hundreds or thousands, drastically reduces the computational overhead for experimentation.
• Reduced Computational Costs: Faster inference directly translates to lower operational costs, especially in cloud environments where compute time is a major expense. This democratizes access to powerful generative capabilities, making them more accessible to smaller teams and individual researchers.
• Stable Training: The MMD-based inductive learning objective is a game-changer for training stability. Practitioners often spend considerable time debugging unstable GAN or diffusion distillation training processes. IMM's stable optimization means less time spent on hyperparameter tuning and failure modes, leading to more predictable and robust model development.
• Flexible Deployment: The ability to choose the number of inference steps allows for dynamic trade-offs between speed and quality at deployment time. For applications requiring real-time generation (e.g., interactive content creation), fewer steps can be chosen. For offline, high-fidelity production (e.g., film special effects), more steps can be leveraged. This flexibility empowers deployers to optimize for diverse use cases.
• Parameter Efficiency: IMM's capability to match or exceed the quality of much larger autoregressive models (e.g., 600M IMM vs. 2B VAR) suggests a higher parameter efficiency. This means smaller models can achieve comparable performance, further reducing memory footprint and deployment costs, especially on edge devices or in resource-constrained environments.

For Infrastructure Teams:

• Optimized Resource Utilization: Faster inference means GPUs are utilized more efficiently, completing tasks in less time. This can lead to better scheduling, higher throughput, and potentially the need for fewer computational resources for a given workload.
• Simplified Deployment Pipelines: Stable training reduces the complexity of monitoring and managing training jobs. Predictable convergence means fewer retraining cycles due to instability, simplifying MLOps pipelines.

Tradeoffs and Limitations:

While IMM presents significant advantages, it's important to consider potential tradeoffs:

• Kernel Choice for MMD: The performance of MMD can be sensitive to the choice of kernel function and its hyperparameters. While standard kernels are mentioned, fine-tuning these for specific datasets might still require expertise.
• Inductive Assumption Validity: The success of inductive learning relies on the assumption that the intermediate model distribution PSR is sufficiently close to the ground truth QS. While empirically shown to be effective, there might be scenarios or data distributions where this assumption holds less strongly, potentially affecting performance or requiring more intermediate steps.
• Complexity of Learning Large Jumps: While the network has the capacity to learn large jumps by injecting S, the complexity of the function it needs to approximate can still be substantial. The training process must effectively learn this highly non-linear mapping, which might require significant data and model capacity.
• Comparison to Future Baselines: The field of generative AI is moving rapidly. While IMM outperforms current diffusion and autoregressive baselines, continuous research in areas like latent diffusion or new distillation techniques might introduce new state-of-the-art models that IMM would need to be re-evaluated against.

In conclusion, IMM offers a compelling solution to the generative trilemma, providing a practical pathway to build and deploy high-quality generative models with unprecedented speed and stability. Its implications extend to making advanced generative AI more accessible, cost-effective, and versatile for a wide range of real-world applications.

Key Takeaways

• Generative Trilemma Solution: Inductive Moment Matching (IMM) directly addresses the challenge of simultaneously achieving high sample quality, stable training, and efficient inference in generative models, a problem that plagues existing methods like diffusion models and GANs.
• Enhanced Inference Efficiency: IMM significantly speeds up inference by injecting the target time step S into the neural network, enabling the model to learn complex, non-linear, large-jump mappings across the generative process, reducing steps from hundreds to as few as 8-16.
• Stable Training with MMD: The training objective leverages Maximum Mean Discrepancy (MMD) for stable distribution matching, avoiding the adversarial instabilities of GANs, and implicitly matching all moments of the distributions.
• Inductive Learning for Easier Optimization: IMM incorporates an inductive learning approach where the model's own intermediate distribution serves as a training target, simplifying the learning task and requiring fewer samples for accurate moment estimation compared to direct matching of distant distributions.
• Superior Performance: IMM achieves an FID of 1.99 with 8 steps and 1.90 with 16 steps, outperforming diffusion and flow matching baselines with the same architecture, and matching or exceeding autoregressive models like VAR with significantly fewer parameters (e.g., 600M IMM vs. 2B VAR).
• Scalability and Flexibility: The method scales well with increasing computational resources and offers flexible control over the number of inference steps, allowing practitioners to balance speed and quality based on specific application requirements.

About the Speaker(s)

Linqi (Alex) Zhou is the primary presenter of this work, indicating his significant role in the research and development of Inductive Moment Matching. He is affiliated with Stanford University, where he collaborates with his advisor.

Stefano Ermon is mentioned as Alex Zhou's advisor at Stanford, highlighting his role as a senior researcher and academic mentor in this project. His expertise in machine learning likely guided the theoretical foundations and experimental design of IMM.

Jiaming Song is identified as the Chief Scientist at Luma AI. His involvement underscores the practical relevance and potential industry applications of IMM, particularly in the context of large-scale foundation models for visual generation, such as those Luma AI focuses on for realistic text-to-video capabilities.