Generalization Dynamics of Linear Diffusion Models

TL;DR

This work probes how finite data and data-structure shape the generalization of diffusion-based generative models using a linear, Gaussian-parameterization framework. By linking the covariance spectrum, particularly power-law hierarchies, to learning dynamics and generalization via replica theory, it identifies two regimes: when $N<d$ hierarchical spectra and regularization mitigate overfitting, and when $N>d$ the generated distribution approaches the population optimum with a $DKL$ scaling of $\sim d/(4N)$, largely independent of the distribution. It also contrasts linear and nonlinear models, showing similar qualitative behavior but notable differences in leading directions, and demonstrates the beneficial role of regularization and early stopping. Overall, the paper clarifies how sample complexity and data hierarchy govern generalization in diffusion models, informing regularization strategies and training protocols for finite-data regimes.

Abstract

Diffusion models are powerful generative models that produce high-quality samples from complex data. While their infinite-data behavior is well understood, their generalization with finite data remains less clear. Classical learning theory predicts that generalization occurs at a sample complexity that is exponential in the dimension, far exceeding practical needs. We address this gap by analyzing diffusion models through the lens of data covariance spectra, which often follow power-law decays, reflecting the hierarchical structure of real data. To understand whether such a hierarchical structure can benefit learning in diffusion models, we develop a theoretical framework based on linear neural networks, congruent with a Gaussian hypothesis on the data. We quantify how the hierarchical organization of variance in the data and regularization impacts generalization. We find two regimes: When $N <d$, not all directions of variation are present in the training data, which results in a large gap between training and test loss. In this regime, we demonstrate how a strongly hierarchical data structure, as well as regularization and early stopping help to prevent overfitting. For $N > d$, we find that the sampling distributions of linear diffusion models approach their optimum (measured by the Kullback-Leibler divergence) linearly with $d/N$, independent of the specifics of the data distribution. Our work clarifies how sample complexity governs generalization in a simple model of diffusion-based generative models.

Figures (10)

• Figure 1: a) Eigenvalues of covariances obtained from image data sets, sorted by rank. b) - d) Example images from image datasets. e) - g) Top three leading eigenvectors of covariance matrices obtained from the full image dataset. h) Prediction for test and training loss of linear diffusion models trained on $N$ samples $d=100$-dimensional samples from $\mathcal{N}(0,\Sigma)$, where the eigenvalues of $\Sigma$ follow a powerlaw with exponent $k$ normed such that $\Tr \Sigma/d =1$. i) - k) Test and train loss of trained diffusion models with linear and U-net architecture trained on $N$ training data. Test losses are averaged over $10^4$ samples from the test set. Training losses are computed using at $\max(N,10^4)$ training data. Grey lines show prediction from replica theory. l) Kullback-Leibler divergence between sample distribution of linear diffusion models with regularization $c=10^{-4}$ and $\mathcal{N}(0,\Sigma)$, where the eigenvalues of $\Sigma$ follow the same powerlaw as in h). Symbols are averages over $10$ random draws of the training sets, error bars report one standard deviation, but are typically smaller than the symbol size. m) - o) are equivalent to l), but for $\Sigma$ originating from the CelebA, MNIST, and CIFAR-10 datasets, respectively. Grey lines show prediction \ref{['eq:dkl_from_replica']} from replica theory.
• Figure 2: a)-b) Kullback-Leibler divergence between sample distribution of linear diffusion models from Gaussian ground truth for varying levels of regularization number of data drawn from ${\mathcal{N}\left(0,\text{Id}\right)}$. c)-f) are equivalent to a), b) but for data drawn from and $\mathcal{N}(0,\Sigma)$, where the eigenvalues of $\Sigma$ follow a powerlaw with $\lambda_{\nu}\sim \nu^{-k}$. Symbols are averages over $10$ random draws of the training sets, error bars report one standard deviation, but are typically smaller than the symbol size. Lines show prediction \ref{['eq:dkl_from_replica']} from replica theory.
• Figure 3: a) Test loss of linear diffusion models trained on $N=0.6d$ samples from a centered Gaussian with powerlaw eigenvalues, $\lambda_{\nu}\sim \nu^{-k}$. b)-d) Are equivalent to b), but for fixed level of noise, $t$. e)-f) are equivalent to b), but for linear denoisers trained with increasing regularization strength $c$. g) Training time $\tau$ with optimal test loss as a function of $k$ for different fractions $N/d$. h) same as g), but for different fractions $N/d$. All curves are averages over $5$ draws of the data and $d=10^3$.
• Figure 4: a)-c) Standard test loss of linear diffusion models \ref{['eq:loss']}, trained on increasing fractions of data as a function of training steps $\tau$ times learning rate $\eta$. d)-e) Same as a)-c), but for test loss corresponding to predicting the data instead of the noise, \ref{['eq:loss_on_data']}. Data are drawn from a centered Gaussian with powerlaw eigenvalues, $\lambda_{\nu}\sim \nu^{-k}$. b)-c) All curves are averages over $5$ draws of the data and $d=10^3$.
• Figure 5: Left column: a) Difference of denoisers from reference model, trained on increasing numbers of data, averaged over $100$ test data points. Blue squares are linear models, pink diamonds are U-nets. Blue lines show prediction from replica calculation. b) Similarity between samples generated from U-net architecture diffusion models and closest training data point, averaged over $400$ generated data points. c)-d) are equivalent to a)-b) but for CIFAR-10 data. Right column: comparison of generated images vs. closest training example for models trained on increasing number of data.