A theory of learning data statistics in diffusion models, from easy to hard
Abstract
While diffusion models have emerged as a powerful class of generative models, their learning dynamics remain poorly understood. We address this issue first by empirically showing that standard diffusion models trained on natural images exhibit a simplicity bias, learning simple, pair-wise input statistics before specializing to higher-order correlations. We reproduce this behaviour in simple denoisers trained on a minimal data model, the mixed cumulant model, where we precisely control both pair-wise and higher-order correlations of the inputs. We identify a scalar invariant of the model that governs the sample complexity of learning pair-wise and higher-order correlations that we call the diffusion information exponent, in analogy to related invariants in different learning paradigms. Using this invariant, we prove that the denoiser learns simple, pair-wise statistics of the inputs at linear sample complexity, while more complex higher-order statistics, such as the fourth cumulant, require at least cubic sample complexity. We also prove that the sample complexity of learning the fourth cumulant is linear if pair-wise and higher-order statistics share a correlated latent structure. Our work describes a key mechanism for how diffusion models can learn distributions of increasing complexity.
Lay Summary
Diffusion models are the technology behind modern AI image generators, yet we still do not understand how they learn to produce such realistic pictures. Which features of the data do they pick up, and in what order? We trained image generators on standard datasets and tested them throughout training, comparing their performance on real images against simplified versions that retained only basic properties such as average brightness and pixel correlations. We found that the model first learns these simple statistics. The complex structure that makes images look truly realistic, like edges and fine textures, is learnt only later. To explain this, we built a mathematical model in which both simple and complex structure can be controlled exactly. We proved that one quantity, which we call the "diffusion information exponent", determines how much data is needed to learn each level of complexity. We also showed that when different features of the data are tied together, as they typically are in real images, complex features are picked up much faster than the worst case suggests. This work clarifies a fundamental aspect of how image generators succeed, and offers a principled handle for designing more efficient training methods.