Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models

1. Problem Statement

Current state-of-the-art generative models, specifically Diffusion Models and Flow Matching models, face a critical computational bottleneck: Log-Likelihood Evaluation.

• The Paradox: While these models excel at generating high-quality samples, computing their exact log-likelihood (log-density) is prohibitively expensive. It typically requires integrating the divergence of the velocity field along the entire trajectory using numerical solvers, necessitating hundreds to thousands of Neural Function Evaluations (NFEs).
• The Gap: Recent advancements in distillation methods (e.g., Consistency Models, Shortcut Models, MeanFlow) have successfully accelerated sampling to just a few steps (e.g., 1–10 NFEs). However, these methods generally abandon the ability to compute likelihoods or still require expensive full-trajectory integration for likelihood evaluation.
• Consequence: This limitation prevents the use of likelihood-based objectives (e.g., Maximum Likelihood Estimation, Reinforcement Learning with PPO/DPO, model selection) in fast few-step flow-based models.

2. Methodology: Fast Flow Joint Distillation (F2D2)

The authors propose F2D2, a framework that simultaneously accelerates both sampling and likelihood evaluation by jointly distilling the sampling trajectory and the cumulative divergence into a single model.

Core Insight

In Continuous Normalizing Flows (CNFs) and Flow Matching, the sampling trajectory ($x_t$) and the log-density evolution ($z_t = \log p_t(x_t)$) are governed by coupled Ordinary Differential Equations (ODEs) driven by the same underlying velocity field $v_\theta$:

1. Sampling: $\frac{d}{dt}x_t = v_\theta(x_t, t)$
2. Likelihood: $\frac{d}{dt} \log p_t(x_t) = -\text{div}(v_\theta(x_t, t))$

Since both depend on the same velocity field, the authors argue that a single model can learn to predict the flow map for the joint state $(x_t, z_t)$, effectively learning to skip steps for both the trajectory and the divergence accumulation simultaneously.

Technical Formulation

• Joint Flow Map: The model parametrizes a joint flow map $\Phi_Y$ that maps a state at time $t$ to time $s$:
  $$\Phi_Y(y_t, t, s) = y_t + (s-t)f_\theta(x_t, t, s)$$
  where $y_t = (x_t, z_t)^\top$ and $f_\theta$ outputs both the average velocity $u_\theta$ and the average divergence $D_\theta$.
• Architecture: The model uses a shared backbone (e.g., U-Net) with two separate prediction heads:
  1. A vector head for the velocity $u_\theta$.
  2. A scalar head for the divergence $D_\theta$.
• Training Objectives: The framework employs a joint self-distillation objective consisting of four loss terms:
  1. Velocity Matching ($L_{VM}$): Ensures the instantaneous velocity matches the ground truth (Flow Matching loss).
  2. Flow Map Consistency ($L_u$): Enforces the flow map property (e.g., Semigroup property for Shortcut models or Eulerian identity for MeanFlow) on the velocity.
  3. Divergence Matching ($L_{div}$): Ensures the instantaneous divergence prediction matches the true divergence (estimated via Hutchinson trace estimator).
  4. Divergence Consistency ($L_D$): Enforces the flow map property on the divergence component, ensuring the accumulated log-likelihood is consistent with the trajectory.

Instantiations

F2D2 is modular and compatible with existing few-step methods:

• Shortcut-F2D2: Based on the Semigroup property (consistency across time intervals).
• MeanFlow-F2D2: Based on the Eulerian equation (time-averaged velocity identity).
• LSD-F2D2: Based on Lagrangian self-distillation.

Practical Design Choices

• Hutchinson Trace Estimator: Used to approximate the divergence $\text{div}(v)$ efficiently ($O(1)$ cost) instead of exact $O(d)$ computation.
• Staged Training: The model is often warm-started with a pre-trained sampling-only teacher to ensure stable velocity predictions before introducing the divergence head.
• Forward-Only Compatibility: The method adapts to existing "forward-only" (t $\le$ s) pretrained models by using first-order approximations for backward integration during likelihood evaluation.

3. Key Contributions

1. First Joint Distillation Framework: F2D2 is the first method to enable accurate few-step log-likelihood evaluation in diffusion/flow-based models without sacrificing sample quality.
2. Modularity: It acts as a "plug-and-play" module that can be added to existing few-step models (Shortcut, MeanFlow) by simply adding a divergence prediction head and modifying the loss function.
3. New Application (Self-Guidance): The authors introduce Maximum Likelihood Self-Guidance, a test-time intervention where the initial noise $x_0$ is optimized using the model's own likelihood predictions to improve sample quality.
4. Theoretical Justification: Provides rigorous proofs characterizing valid joint flow maps under Lagrangian, Eulerian, and Semigroup conditions.

4. Experimental Results

Experiments were conducted on CIFAR-10, ImageNet 64x64, and CelebA-64.

• Likelihood Accuracy:
  • Standard few-step models (Shortcut, MeanFlow) produce invalid or highly inaccurate NLL (Negative Log-Likelihood) estimates (e.g., -100s BPD) when using few steps.
  • F2D2 achieves calibrated NLL close to the teacher model (trained with 1024 steps). For example, on CIFAR-10, MeanFlow-F2D2 achieves ~1.63 BPD with 2 steps, matching the teacher's 3.12 BPD (where lower is better, and the teacher is the ground truth reference).
• Sample Quality (FID):
  • F2D2 maintains competitive or improved FID scores compared to the base few-step models.
  • Key Result: A 2-step MeanFlow-F2D2 with Self-Guidance achieves a lower FID than a 1024-step Flow Matching model of the same size on CIFAR-10. This demonstrates that fast likelihood evaluation enables superior generation via optimization.
• 2D Synthetic Data: On a 2D checkerboard dataset with known ground truth, F2D2 accurately recovers the target density distribution even with 1 NFE, whereas baseline methods fail catastrophically.

5. Significance and Impact

• Solving a Long-Standing Bottleneck: F2D2 removes the computational barrier preventing the use of likelihood-based objectives in modern, fast generative models.
• Enabling New Training Paradigms: By making likelihoods tractable in few steps, F2D2 unlocks applications such as:
  • Reinforcement Learning (RL): Using likelihoods as rewards for PPO/DPO.
  • Model Selection: Comparing models based on exact likelihood rather than just FID.
  • Self-Guidance: Optimizing generation quality without external reward models.
• Efficiency: It reduces the NFEs for likelihood evaluation by two orders of magnitude (from ~1000 to ~1-10), making these models viable for real-time or resource-constrained applications requiring both generation and analysis.

In summary, F2D2 bridges the gap between the speed of distillation-based sampling and the analytical rigor of likelihood-based inference, establishing a new standard for efficient and versatile flow-based generative modeling.