Generalization Dynamics of Linear Diffusion Models

TL;DR

This paper analyzes how the hierarchical structure of data affects the generalization of linear diffusion models, revealing regimes where regularization and early stopping improve performance, and quantifying sample complexity effects.

Contribution

It introduces a theoretical framework based on linear neural networks and data covariance spectra to understand diffusion model generalization with finite data.

Findings

Hierarchical data structure and regularization mitigate overfitting in low-sample regimes.

For large sample sizes, the divergence approaches its optimum linearly with d/N.

The analysis clarifies the role of data complexity in diffusion model generalization.

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.