Understanding Generalization in Diffusion Distillation via Probability Flow Distance

1. Problem Statement

Diffusion distillation is a critical technique for deploying diffusion models efficiently, enabling lightweight, few-step generation. However, evaluating the generalization ability of distilled models remains a significant challenge due to the lack of metrics that are both theoretically rigorous and computationally practical.

• Limitations of Existing Metrics:
  • Practical Metrics (e.g., FID, IS): Focus on generation quality but fail to distinguish between memorization (copying training data) and generalization (learning the underlying data distribution). Both can yield high-quality outputs.
  • Theoretical Metrics (e.g., KL Divergence, Wasserstein Distance): While theoretically sound, they are often computationally intractable for high-dimensional data (e.g., images) or require prohibitive sample sizes for accurate estimation.
  • Sample-level Metrics: Approaches measuring the distance to the nearest neighbor in the training set can be misleading (e.g., misclassifying pure noise as generalized output).

The authors aim to bridge this gap by introducing a metric that can quantitatively measure the transition from memorization to generalization in diffusion distillation.

2. Methodology: Probability Flow Distance (PFD)

The core contribution is the Probability Flow Distance (PFD), a novel metric grounded in the Probability Flow Ordinary Differential Equation (PF-ODE).

• Theoretical Foundation:

  • Diffusion models utilize a forward process (adding noise) and a backward PF-ODE (removing noise to generate data).
  • The PF-ODE induces a unique, deterministic mapping $\Phi_p$ from a standard Gaussian noise distribution ($\mathcal{N}(0, T^2\mathbf{I})$) to the data distribution $p$.
  • PFD measures the distance between two distributions $p$ and $q$ by comparing their respective noise-to-data mappings ($\Phi_p$ and $\Phi_q$) starting from the same noise input.

• Definition:
  $$\text{PFD}(p, q) = \left( \mathbb{E}_{\mathbf{x}_T} \left[ \| \Psi \circ \Phi_p(\mathbf{x}_T) - \Psi \circ \Phi_q(\mathbf{x}_T) \|_2^2 \right] \right)^{1/2}$$
  Where:

  • $\mathbf{x}_T$ is the initial noise.
  • $\Phi_p(\mathbf{x}_T)$ is the sample generated by the model $p$ from noise $\mathbf{x}_T$.
  • $\Psi(\cdot)$ is an image descriptor (e.g., DINOv2, SSCD) to capture perceptual quality, though the theoretical analysis assumes an identity mapping.

• Properties:

  • Metric Axioms: PFD satisfies positivity, identity, symmetry, and the triangle inequality.
  • Efficiency: Unlike Wasserstein distance ($O(M^2)$ complexity) or KL divergence (requiring density estimation), PFD has $O(M)$ computational complexity and requires significantly fewer samples for accurate estimation.
  • Theoretical Guarantee: Under Lipschitz continuity assumptions of score functions, the empirical estimate of PFD converges to the true value with high probability.

3. Key Contributions

1. Proposal of PFD: A theoretically grounded, computationally efficient metric that explicitly quantifies generalization in diffusion models, distinguishing it from memorization.
2. Quantitative Scaling Laws: The authors identify a universal scaling behavior governing the transition from memorization to generalization, governed by the ratio $N / \sqrt{|\theta|}$, where $N$ is the dataset size and $|\theta|$ is the number of model parameters.
3. Learning Dynamics Analysis:
   • Epoch-wise Double Descent: In the generalization regime (sufficient data), the generalization error exhibits a "double descent" curve (decrease $\to$ increase $\to$ decrease) during training, a phenomenon previously observed in supervised learning but unverified in diffusion distillation.
   • Early Learning: In regimes with limited data, models exhibit "early learning," where they generalize initially but later transition to memorization as training progresses.
4. Bias-Variance Decomposition: The paper extends classical statistical learning theory to diffusion models, decomposing generalization error into bias and variance components, revealing a U-shaped error curve consistent with supervised learning.

4. Experimental Results

The authors evaluated PFD on CIFAR-10 and ImageNet datasets using various U-Net architectures (ranging from 0.9M to 295.9M parameters) and distillation sample sizes ($N$ from $2^6$ to $2^{16}$).

• Superiority over Existing Metrics:

  • In the "Memorization-to-Generalization" (MtoG) transition, traditional metrics like FID, IS, and training loss failed to track the transition accurately. They often degraded due to visual quality drops in the transition phase or were insensitive to distributional shifts.
  • PFD was the only metric that consistently tracked the transition: Memorization error ($\mathcal{E}_{mem}$) increased with $N$, while Generalization error ($\mathcal{E}_{gen}$) decreased, clearly delineating the regimes.

• Scaling Behavior:

  • When plotting error against $\log_2(N)$, curves for different model sizes were distinct.
  • When plotting against $\log_2(N / \sqrt{|\theta|})$, all curves collapsed into a single universal curve, confirming the scaling law.

• Learning Dynamics:

  • Double Descent: For large $N$ ($2^{16}$), $\mathcal{E}_{gen}$ dropped, rose (intermediate epochs), and dropped again, while standard loss functions decreased monotonically.
  • Early Learning: For small $N$ ($2^6$), $\mathcal{E}_{gen}$ dropped initially but rose sharply as the model overfit/memorized the data.

• Bias-Variance Trade-off:

  • Increasing model capacity reduced bias but increased variance.
  • Increasing dataset size reduced both bias and variance.

5. Significance and Impact

• Theoretical Bridge: PFD provides the first rigorous framework to apply statistical learning theory (bias-variance, scaling laws) to diffusion models, a field previously dominated by empirical observations.
• Practical Guidance: The scaling law ($N / \sqrt{|\theta|}$) offers actionable guidance for practitioners on how to balance model size and data volume to achieve optimal generalization without unnecessary over-parameterization.
• Safety and Copyright: By distinguishing memorization from generalization, PFD helps mitigate risks related to privacy leakage and copyright infringement, as it can detect when a model is merely copying training data rather than learning the underlying distribution.
• Beyond Distillation: While designed for distillation, the authors argue that PFD is applicable to training diffusion models from scratch, especially given that modern large-scale datasets often contain synthetic data, effectively creating a teacher-student dynamic implicitly.

In summary, this work establishes Probability Flow Distance as a foundational tool for understanding, evaluating, and optimizing the generalization capabilities of diffusion models, moving the field from qualitative assessment to quantitative, theory-driven analysis.