The Unreasonable Effectiveness of Gaussian Score Approximation for Diffusion Models and its Applications

TL;DR

This work investigates why diffusion models perform so well by examining the geometry of the learned score function. It shows that, at high noise, the neural score is dominated by a Gaussian (linear) approximation of the data, and at lower noise levels, Gaussian mixtures better capture the learned structure, while the exact training-score (delta) model underfits generalization. The authors derive exact Gaussian PF-ODE solutions, reveal a rotation-like geometry of sampling trajectories, and identify a feature-emergence order governed by data variance. They introduce analytical teleportation to accelerate sampling, demonstrating substantial reductions in neural function evaluations with minimal quality loss, and propose design implications for training and architecture to leverage the far-field Gaussian structure for diffusion models.

Abstract

By learning the gradient of smoothed data distributions, diffusion models can iteratively generate samples from complex distributions. The learned score function enables their generalization capabilities, but how the learned score relates to the score of the underlying data manifold remains largely unclear. Here, we aim to elucidate this relationship by comparing learned neural scores to the scores of two kinds of analytically tractable distributions: Gaussians and Gaussian mixtures. The simplicity of the Gaussian model makes it theoretically attractive, and we show that it admits a closed-form solution and predicts many qualitative aspects of sample generation dynamics. We claim that the learned neural score is dominated by its linear (Gaussian) approximation for moderate to high noise scales, and supply both theoretical and empirical arguments to support this claim. Moreover, the Gaussian approximation empirically works for a larger range of noise scales than naive theory suggests it should, and is preferentially learned early in training. At smaller noise scales, we observe that learned scores are better described by a coarse-grained (Gaussian mixture) approximation of training data than by the score of the training distribution, a finding consistent with generalization. Our findings enable us to precisely predict the initial phase of trained models' sampling trajectories through their Gaussian approximations. We show that this allows the skipping of the first 15-30% of sampling steps while maintaining high sample quality (with a near state-of-the-art FID score of 1.93 on CIFAR-10 unconditional generation). This forms the foundation of a novel hybrid sampling method, termed analytical teleportation, which can seamlessly integrate with and accelerate existing samplers, including DPM-Solver-v3 and UniPC. Our findings suggest ways to improve the design and training of diffusion models.

Figures (29)

• Figure 1: Gaussian score well-approximates the learned neural score. A. Schematic illustrating our work's main claim. In the high noise regime, the neural score is well-approximated by both the Gaussian/linear score and the score of the training set; in the low noise regime, neural scores are better described by the Gaussian model. B. Visual demonstration of the effect. Denoiser outputs along PF-ODE trajectories were similar at high noise, regardless of the score model (neural score, Gaussian score, or delta mixture score); at small noise scales, neural denoiser outputs more closely resemble those of the Gaussian model.
• Figure 2: Structure of Gaussian, Gaussian mixture, and delta mixture scores.Upper row: Examples of 2D probability density functions for Gaussian, Gaussian mixture model (GMM), and delta mixture distributions. Lower row: The respective score vector fields. For mixture models, the space is colored based on the weighting function $w_i(\mathbf{x})$ from Eq. \ref{['eq:general_gmm_score_main']}. We assign a unique color to each Gaussian component (e.g., RGB colors for three components), and the space is colored according to the mixing of component colors weighted by $w_i(\mathbf{x})$.
• Figure 3: Analytical solution to sample generation dynamics for Gaussian model. A. The noise and signal schedule $\sigma_t^2$ and $\alpha_t^2$ from ddpm-CIFAR-10. B.$\bar{\psi}(t,\lambda)$ governs the dynamics of $\mathbf{x}_t$ along each principal axis $\mathbf{u}_k$. C.$\bar{\xi}(t,\lambda)$ governs the dynamics of the endpoint estimate $\hat{\mathbf{x}}_0(\mathbf{x}_t)$ along each PC, normalized by the standard deviation $\sqrt{\lambda}$. D. Time derivative of $\bar{\xi}(t,\lambda)/\sqrt{\lambda}$, highlighting the 'critical period' when each feature develops the most rapidly. E.$\sqrt{\lambda/(\sigma_{t'}^2+\lambda\alpha_{t'}^2)}$, which quantifies the amplification effect of a perturbation along PC $\mathbf{u}_k$ at time $t'$ (Eq.\ref{['eq:y_perturb_formula']}).
• Figure 4: Geometry of sample generation trajectories for the VP-SDE.
• Figure 5: Qualitative aspects of sample generation consistent with Gaussian theory. A. Trajectories of states $G(\mathbf{x}_t)$ (top row), scaled differences between nearby states $G(k (\mathbf{x}_{t-1} - \mathbf{x}_{t}))$ (middle row), and denoiser / projected outcome $G(\hat{\mathbf{x}}_0(\mathbf{x}_t))$ (bottom row). B. Trajectory of $\mathbf{x}_t$ projected onto the plane spanned by $\mathbf{x}_T$ and $\mathbf{x}_0$. Trajectories are effectively two-dimensional, and resemble a rotation from $\mathbf{x}_T$ to $\mathbf{x}_0$. Note that both feature emergence order and the low dimensionality of trajectories are consistent with the Gaussian model.