Memorization and Generalization in Generative Diffusion under the Manifold Hypothesis

TL;DR

This paper analyzes how diffusion models memorize and generalize structured data on manifolds, revealing that data structure enhances generalization and that memorization and generalization phases are distinct with optimal generalization occurring during memorization.

Contribution

It provides a theoretical framework connecting diffusion models' behavior to the Random Energy Model, characterizing memorization and generalization phases in high-dimensional structured data.

Findings

Memorization time $t_c$ decreases with data structure complexity.

Generalization is optimized at a time $t_g$ before memorization.

Structured data avoids the curse of dimensionality in memorization.

Abstract

We study the memorization and generalization capabilities of Diffusion Models (DMs) when data lies on a structured latent manifold. Specifically, we consider a set of $P$ data points in $N$ dimensions confined to a latent subspace of dimension $D = \alpha_D N$, following the Hidden Manifold Model (HMM). We analyze the reverse diffusion process using the empirical score function as a proxy, and characterize it in the high-dimensional limit $P = \exp(\alpha N)$, $N \gg 1$, by exploiting a connection with the Random Energy Model (REM). We show that a characteristic time $t_o$ marks the emergence of traps in the time-dependent potential, which however do not affect typical trajectories. The size of their basins of attraction is computed at all times. We derive the collapse time $t_c < t_o$, at which trajectories fall into the basin of a training point, signaling memorization. An explicit formula for $t_c$ as a function of $P$ and $\alpha_D$ shows that the curse of dimensionality is avoided for structured data ($\alpha_D \ll 1$), even with nonlinear manifolds. We also prove that collapse corresponds to the condensation transition in the REM. Generalization is quantified via the Kullback-Leibler divergence between the exact distribution and the reverse one at time $t$. We find a distinct time $t_g < t_c < t_o$ minimizing this divergence. Surprisingly, the best generalization occurs inside the memorization phase. We conclude that generalization in DMs improves with data structure, as $t_g \to 0$ faster than $t_c$ when $\alpha_D \to 0$.