From Easy to Hard: A Theory of How Diffusion Models Learn Data Statistics
A theory of learning data statistics in diffusion models, from easy to hard
This paper introduces a theoretical framework to understand the learning dynamics of diffusion models, revealing a "distributional simplicity bias." Using a "mixed cumulant model," the authors prove that diffusion models learn lower-order statistics (like mean and covariance) at linear sample complexity before capturing higher-order correlations (like 4th-order cumulants), which require at least cubic sample complexity.
TL;DR
Why do diffusion models first generate blurry blobs before fine-grained details? This paper provides the mathematical answer: Distributional Simplicity Bias. By analyzing a "mixed cumulant model," the researchers prove that diffusion models learn pairwise statistics (covariance) drastically faster than higher-order correlations. They identify a critical metric called the Diffusion Information exponent () that dictates the "wait time" (sample complexity) for a model to "see" specific data features.
The Mystery of Sequential Learning
When training a U-Net on CIFAR-10, researchers noticed a curious trend: for the first thousand steps, the model's loss on real images is identical to its loss on "Gaussian clones" (images with the same mean and covariance but no higher-order structure).
The model essentially treats the data as a simple Gaussian distribution before it "awakens" to the complex textures and shapes that define real objects. This work asks: Why does this happen, and what determines the speed of this transition?
Figure 1: As training progresses, the generated samples evolve from noise to simple blobs (mean/covariance level) to structured images, reflecting the internal learning of higher-order statistics.
The Mixed Cumulant Model (MCM)
To solve this, the authors used a controlled "Spiked" data model. Imagine data that looks like isotropic noise except along two hidden directions (spikes):
- The Covariance Spike (): Hidden in the 2nd-order statistics (easy).
- The Cumulant Spike (): Hidden only in the 4th-order statistics (hard).
The core of their theory is the Diffusion Information Exponent (). By expanding the denoiser’s nonlinearity and the data's likelihood ratio into Hermite polynomials, they found that the model only begins to learn a feature of order when it has seen roughly samples.
Key Insights: Why Your Optimization Might Fail
One of the paper's most startling revelations concerns Unconstrained SGD. In many theoretical setups, researchers ignore the norm of the weights. However, this paper proves that for many common activation functions (like the "optimal" denoiser derived from the data itself), unconstrained SGD possesses an attracting fixed point at zero.
Basically, the "noise" in the diffusion process creates a radial contraction force. If your model isn't constrained to a sphere or isn't "wide" enough (over-parameterized), the weights will simply shrink to zero, and the model will learn nothing.
Figure 2: Different activation functions () show varying levels of "contraction." If , the model is sucked into a trivial solution at the origin.
Solving the "Hard" Part: Depth and Correlations
How do real models escape this? The authors highlight two escape hatches:
- Correlated Latents: If the "simple" Gaussian features and "hard" non-linear features are statistically correlated, the model can "piggyback" on the easy features to learn the hard ones in linear time.
- Architecture Depth: While a single-neuron model might get stuck, deeper and wider networks (ResNets/Autoencoders) can re-parameterize the input, effectively flipping the sign of the contraction and making hard features learnable.
Figure 3: Depth and Width help. Even when the mathematical "contraction" suggests failure, deeper networks eventually "find" the hidden cumulant spike.
Conclusion
This work moves diffusion theory away from "everything is Gaussian" toward a rigorous understanding of feature hierarchy. It defines a mathematical "speed limit" for learning data complexity and provides a theoretical justification for why we need the massive over-parameterization found in modern U-Nets and Transformers: it isn't just for memory, it's to break the "simplicity bias" and capture the complex correlations that make pixels into people.
