Generalization Dynamics of Linear Diffusion Models

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.

💡 Research Summary

The paper investigates the generalization behavior of diffusion‑based generative models in the regime of finite training data, using a tractable linear‑network setting. Assuming the data distribution is a multivariate Gaussian with covariance Σ, the authors focus on the spectrum of Σ, which in many real‑world image datasets follows a power‑law λₙ ∝ n^{‑k}. The exponent k quantifies hierarchical structure: larger k means a few leading eigen‑directions dominate the variance.

A linear denoiser is trained at each diffusion timestep with an L₂ regularizer γₜ. The learning outcome depends only on the empirical mean μ̂ and empirical covariance Σ̂ computed from N training samples. When N < d (the dimensionality), Σ̂ has a nullspace of dimension at least d‑N, causing many eigenvalues to be exactly zero. Equation (4) shows that the gap between training loss R and test loss L_test is amplified for directions with small eigenvalues and early timesteps, leading to severe over‑fitting.

The authors demonstrate analytically that a hierarchical spectrum mitigates this problem. Because the leading eigenvalues are large, they can be estimated accurately even with few samples, while the missing directions correspond to small variance in the true data, so their absence in Σ̂ does not hurt performance much. Regularization (γₜ) adds a term to the denominator of (4), reducing the impact of the nullspace; the optimal regularization strength decreases with both N and k. Early stopping also benefits from the slower learning of low‑variance directions, providing a wider window to stop before over‑fitting occurs.

To quantify the overall distributional error, the Kullback‑Leibler divergence D_KL(ρ_N‖ρ) between the generated Gaussian ρ_N = N(μ̂, Σ̂ + cI) and the true data distribution is analyzed using the replica method. The average D_KL over training‑set draws is given by equations (7) and (8), involving a scalar q that depends on the eigenvalues of Σ, the sample size N, and the regularization constant c. Smaller q leads to smaller D_KL. The authors prove an upper bound q ≤ (d/(N·\barλ + c))·Tr