First Learn, Then Memorize: The Spectral Bias of Diffusion Models

TL;DR. When you train a diffusion model on a finite dataset, something odd and well-documented happens: it first learns to generate convincing new samples, and only much later collapses into regurgitating its training set. This paper pins down why, with a mechanism you can state in one sentence: the score-matching loss shows the network m noisy copies of each clean image, and this clustering of the training inputs splits the spectrum of the NTK Gram matrix into two well-separated bulks — large eigenvalues that encode the data distribution's global structure, and tiny eigenvalues (absent when m=1) whose eigenvectors point along the sample-specific noise directions. Since gradient flow learns spectral modes in decreasing eigenvalue order, the gap between the bulks is the window between generalization and memorization, and it widens linearly with dataset size. The authors prove this in the lazy/high-dimensional limit, then show the same two-bulk structure in convolutional NTKs on CelebA and in finite-width U-Nets, and — the best part — intervene on it causally: truncating the Gram matrix at rank r dials the generalization–memorization transition, and an L_2 penalty tuned to kill only the second bulk suppresses memorization entirely.

The puzzle: the empirical score wants to memorize

A diffusion model learns a vector field — the score \mathbf{s}(\mathbf{x}, t) = \nabla_\mathbf{x} \log P_t(\mathbf{x}), the gradient of the log-density of noised data — and generates samples by integrating a reverse-time SDE guided by it. In practice you train by denoising score matching: take a clean sample \mathbf{x}^\nu, noise it to \mathbf{Y}^{\nu\alpha} = e^{-t}\mathbf{x}^\nu + \sqrt{\Delta_t}\,\boldsymbol{\xi}^{\nu\alpha} (where \Delta_t = 1 - e^{-2t} is the noise variance and \boldsymbol{\xi}^{\nu\alpha} is Gaussian), and regress the network's output onto the noise.

Here's the problem. The global minimizer of this loss on a finite dataset of n points is the empirical score — the score of the noised empirical distribution, a mixture of n Gaussians centered on the training points. Integrating the reverse process under the empirical score drives every trajectory back onto a training sample. Exactly. Unless n grows exponentially with the data's intrinsic dimension, perfect optimization means perfect memorization.

And yet diffusion models generalize. The resolution, identified in recent work, is dynamical: training passes through a long window where the model has learned the distribution's structure but hasn't yet memorized, so early stopping gives you a generative model rather than a lookup table. What was missing was the mechanism — what object in the training dynamics creates two timescales, \tau_{\mathrm{gen}} \ll \tau_{\mathrm{mem}}, and why does the gap between them grow with n?

The reframe: score matching is regression on clustered data

The authors' move is to notice that denoising score matching is a non-standard regression problem in a specific, consequential way. In ordinary supervised learning, each training input is an independent draw. In score matching, each clean sample \mathbf{x}^\nu appears m times, with m different noise realizations — in practice, m is roughly the number of epochs, since each pass draws fresh noise. So the N = nm training inputs are not i.i.d.: they form n tight clusters of m points, each cluster huddled around one clean sample.

This parameter m interpolates between two worlds. At m=1 you have standard kernel regression with noisy labels. As m \to \infty you recover the exact DSM loss, whose minimizer is the memorizing empirical score. The interesting physics lives at finite m > 1.

Why does clustering matter? Under gradient flow in the lazy (NTK) regime, the trained score has a closed form governed entirely by the N \times N Gram matrix \mathbf{G}^{\nu\alpha,\mu\beta} = K(\mathbf{Y}^{\nu\alpha}, \mathbf{Y}^{\mu\beta}) of the neural tangent kernel K on the noisy training points:

The exponential filter \mathbf{I} - e^{-\frac{d\tau}{nm}\mathbf{G}} switches on spectral modes in decreasing eigenvalue order: a mode with eigenvalue \lambda is learned at time \tau_\lambda \sim nm/(d\lambda). So the whole question — what is learned when — reduces to: what does the spectrum of \mathbf{G} look like, and what do its eigenvectors encode?

The spectrum: two bulks, one gap

In the high-dimensional limit (n, d \to \infty with \psi_n = n/d fixed), the nonlinear Gram matrix is spectrally equivalent to a linear one:

where the \mu's are constants built from the kernel. The term that has no analogue in classical kernel theory is \mu_B \mathbf{B}_m: a block matrix that is 1 exactly on pairs of points within the same cluster. It exists only because of the repeated noising — at m=1 it degenerates into the identity. This single term restructures the spectrum (their Theorem 2.1) into:

• a generalization bulk of d eigenvalues at \lambda_{\mathrm{gen}} \simeq \mu_1 \psi_n m \lambda_t = \Theta(nm/d) — one per eigenvalue \lambda of the data covariance \boldsymbol{\Sigma} (with \lambda_t = e^{-2t}\lambda + \Delta_t), and
• a memorization bulk of d eigenvalues at \lambda_{\mathrm{mem}} \simeq \mu_I + \mu_B(m-1)\Delta_t/\lambda_t = \Theta(m), which simply does not exist at m=1,

plus two inert atoms that provably contribute nothing to the learned score. The eigenvectors make the labels earn their names. The generalization modes are \mathbf{u}_1^{\nu\alpha} \propto \mathbf{v}_\lambda^\top \mathbf{Y}^{\nu\alpha} — projections of the training points on the principal directions of the population covariance: exactly what a standard kernel learns. The memorization modes contain \boldsymbol{\xi}_\perp^{\nu\alpha} = \boldsymbol{\xi}^{\nu\alpha} - \frac{1}{m}\sum_\beta \boldsymbol{\xi}^{\nu\beta} — the component of each noise realization that distinguishes it from its siblings in the same cluster. These modes resolve the fine-grained irregularity of the empirical score around each individual training point. Learning them is memorization, by construction.

Now the punchline writes itself. The generalization bulk sits at \Theta(nm/d) and is learned at \tau_{\mathrm{gen}} = O(1) in the rescaled time. The memorization bulk sits at \Theta(m), a factor \psi_n = n/d smaller, so it's learned at \tau_{\mathrm{mem}} \sim \psi_n \cdot \tau_{\mathrm{gen}}. Double the dataset and the safe window doubles. This recovers, now with an explicit spectral mechanism, the timescale separation observed empirically by Bonnaire et al. and Favero et al.

The analysis extends beyond the proportional regime. For n \asymp d^k, the generalization bulk splits into sub-bulks learned degree by degree — the classical hierarchy of kernel regression, where low-order moments of the data come first — while the memorization bulk stays put at \Theta(m), independent of n, always last in line.

Bias–variance: why early stopping is not optional

A closed-form bias–variance decomposition of the kernel ridge estimator (Theorem 2.2, derived via replicas) sharpens the story. In the ridgeless limit — i.e., training to convergence — the bias saturates at \Theta(t^{-2}), independent of sample complexity. More data does not help; the memorization modes inject a persistent \Theta(1) error that diverges at small noise levels, which is precisely where memorization lives. At the optimal ridge (equivalently, optimal early stopping), the bias instead falls as \Theta(\psi_n^{-2}t^{-2}). The transition happens exactly when the ridge drops to the O(1) scale of the memorization bulk and stops damping it. This is a crisp statement that early stopping in diffusion training is not a heuristic regularizer — it is the thing standing between you and an irreducible bias contributed by one identifiable spectral component.

From theory to CelebA: making the link causal

The high-dimensional theory assumes inner-product kernels and sub-Gaussian data. Do real images cooperate? The authors compute closed-form convolutional NTK Gram matrices on CelebA (32\times32 grayscale, up to n = 3072, m = 8) and find the same qualitative structure: a first bulk whose eigenvalues scale with n and collapse under rescaling, and a second bulk independent of n. (One honest deviation: the first bulk has \sim n eigenvalues rather than the predicted d, reflecting the power-law covariance of natural images.)

Then come the interventions. Truncation: keep only the top-r eigenmodes of \mathbf{G} and generate with the truncated velocity field — a hard-threshold proxy for early stopping. As r grows, FID falls while memorization fraction f_{\mathrm{mem}} stays near zero; then, at larger r, generated images morph into their paired training images until, at full rank, every sample is a copy. The onset of good generation is essentially independent of n, while the memorization threshold scales linearly with n — all curves collapse when plotted against r/n. That is the theory's prediction, read directly off real images.

Feature learning: they train a finite-width U-Net well outside the lazy regime (n=256, m=8) and track its empirical, time-dependent NTK Gram spectrum. The two-bulk structure with a gap at r \approx n is present at initialization and — despite the kernel visibly drifting — persists across the entire training run, even as f_{\mathrm{mem}} climbs from 0 to 0.86. The gap actually sharpens.

Targeted regularization: finally, they add an L_2 penalty whose strength is set dynamically to the n-th eigenvalue of the empirical Gram matrix — just enough to damp everything below the gap and nothing above it.

Memorization stays exactly zero over 10,000 steps, the test loss doesn't suffer, and the second bulk visibly smooths out during training. This is the difference between correlation and mechanism: they identified the spectral component, aimed a regularizer at it, and memorization switched off.

What this changes, and where to push back

If this picture holds broadly, several things follow. Implicit regularization in diffusion models stops being folklore: the safe training window scales as n/d, and you can in principle monitor it via the empirical NTK spectrum, or enforce it with a spectrum-aware penalty rather than guessing a stopping time. It also reframes what's special about the diffusion loss — the danger isn't the architecture or the data, it's the cluster structure that repeated noising imposes on the regression inputs, a structure that orthodox kernel theory simply has no term for (\mu_B \mathbf{B}_m is the new object).

Caveats worth holding onto. The theorems live in the lazy, infinite-width, fixed-t, frozen-noise regime; real training uses fresh noise every step (m \to \infty in effect), feature learning, and a joint objective across noise levels. The U-Net experiment is encouraging but small (n=256, one architecture, grayscale 32\times32), and the memorization metric (nearest-neighbor ratio with trajectories seeded from training noises) is a deliberately memorization-friendly probe. On natural images the clean two-bulk separation partially merges because of power-law covariance spectra — the paper is upfront that the predicted eigenvalue counts don't match. And the dynamically-tracked eigenvalue regularizer requires repeated Gram eigendecompositions, which won't scale as-is to Stable-Diffusion-sized training. There's also closely related concurrent work (Latourelle-Vigeant et al.) reaching a compatible picture via the m=\infty kernel operator, so credit for the broad phenomenology is shared; this paper's distinctive contributions are the finite-m Gram-matrix mechanism, the bias–variance decomposition, and the causal interventions.

Where to spend your time: Section 2.1 for the linear-equivalent Gram matrix and the eigenvector interpretation (the conceptual core, readable in ten minutes), then Section 3's Figure 4 (the truncation-rank sweep with the r/n collapse) and Figure 5 (the regularized U-Net) for the evidence that the mechanism survives contact with reality. Appendix C.2 is a good short read if you want the relation to the random-features Hessian picture of prior work.