Source-linked AI summary

Why Diffusion Models Don't Memorize: The Role of Implicit Dynamical Regularization in Training

Tony Bonnaire, Raphaël Urfin, Giulio Biroli, Marc Mézard

arXiv:2505.17638v2cs.LGcond-mat.dis-nnstat.ML

TL;DR

The paper asks why diffusion models often generalize despite empirical-score learning eventually causing memorization. Using U-Net experiments and a high-dimensional random-features analysis, it identifies two training timescales: τgen for high-quality generation and τmem for memorization. τgen stays independent of n while τmem grows approximately linearly with n, creating a broad early-stopping window that provides implicit dynamical regularization.

  • Problem

    The paper investigates how diffusion models avoid reproducing training data in practice despite empirical-score learning being associated with eventual memorization.

  • Method

    The authors analyze gradient-descent score-learning dynamics numerically with U-Net models and analytically using a high-dimensional random-features network.

  • Results

    τgen is independent of n, whereas τmem grows approximately linearly with n, producing a broad interval of training times with high-quality, novel samples before memorization.

  • Takeaways & Limitations

    Implicit dynamical regularization can prevent memorization in highly overparameterized diffusion models when training is stopped within the gap between τgen and τmem.

  • Takeaways & Limitations

    The analysis focuses on unconditional diffusion models, and the dependence of the absolute timescales on conditioning remains open.

Abstract

from arXiv · show

Diffusion models have achieved remarkable success across a wide range of generative tasks. A key challenge is understanding the mechanisms that prevent their memorization of training data and allow generalization. In this work, we investigate the role of the training dynamics in the transition from generalization to memorization. Through extensive experiments and theoretical analysis, we identify two distinct timescales: an early time $τ_\mathrm{gen}$ at which models begin to generate high-quality samples, and a later time $τ_\mathrm{mem}$ beyond which memorization emerges. Crucially, we find that $τ_\mathrm{mem}$ increases linearly with the training set size $n$, while $τ_\mathrm{gen}$ remains constant. This creates a growing window of training times with $n$ where models generalize effectively, despite showing strong memorization if training continues beyond it. It is only when $n$ becomes larger than a model-dependent threshold that overfitting disappears at infinite training times. These findings reveal a form of implicit dynamical regularization in the training dynamics, which allow to avoid memorization even in highly overparameterized settings. Our results are supported by numerical experiments with standard U-Net architectures on realistic and synthetic datasets, and by a theoretical analysis using a tractable random features model studied in the high-dimensional limit.

1 Introduction

Diffusion models can reproduce training samples when score learning reaches the empirical score, yet practical training often generalizes instead. The paper investigates how training dynamics drive this transition and identifies distinct generalization and memorization regimes.

  • Setting: generative diffusion and score learning: Diffusion models generate data by reversing a noise-adding forward process with a score function learned through score matching.The score is approximated within a parametrized function class, such as a neural network, using empirical training averages.
  • Contributions and theoretical picture: Two timescales govern expressive models: τgen marks high-quality generation, while later τmem marks increasingly strong reproduction of training samples.The paper studies these timescales as functions of training-set size n and model expressivity p.
  • Contributions and theoretical picture: τgen is independent of n, whereas τmem grows approximately linearly with n, creating an early-stopping window τ ∈ [τgen, τmem] with generalization.The widening gap is attributed to smooth interpolation of an irregular empirical score at low noise, before eventual convergence to the empirical-score minimum.
  • Contributions and theoretical picture: The study combines numerical experiments on U-Net score models and CelebA with high-dimensional analysis of a tractable random-features model.Experiments vary training-set size and model expressivity while measuring sample quality and memorization.
  • Motivation: Exact empirical-score learning eventually reproduces training samples, whereas practical optimization can reach a smoother, training-set-independent score with good generalization.The paper frames understanding this divergence between global minimization and practical training as its central aim.

2 Generalization and memorization during training of diffusion models

Experiments show that diffusion models first improve generation and later memorize, with the memorization transition delayed by larger datasets and affected by model capacity. Training and test losses separate during score overfitting, while full-batch results indicate the delay is not caused by sample repetition.

  • Role of training set size on the learning dynamics: At τ = 100K, generated CelebA samples differ from nearest training neighbors, whereas at about 1.62M steps they are very close to them.The figure compares models trained with n = 1024 at early and late training times.
  • Role of training set size on the learning dynamics: τgen ≈100K updates is independent of n, while fmem remains zero until a later τmem that scales with dataset size.At fixed capacity, FID quickly reaches a minimum near τgen, and normalized memorization curves collapse when plotted against τ/n, increasing around τ/n ≈300.
  • Role of training set size on the learning dynamics: As n increases, the critical overfitting time increases and the training–test loss gap shrinks; for n = 32768, test loss remains close to training loss after 11M steps.Initially, training and test losses are nearly indistinguishable; later, training loss decreases while test loss increases for smaller datasets.
  • Role of training set size on the learning dynamics: Full-batch updates still produce τmem ∝ n, showing that delayed memorization is not explained by unequal sample repetition.The authors attribute the effect to an n-dependent change in the loss landscape.
  • Effect of the model capacity: Larger-capacity U-Nets achieve high-quality generation and begin memorizing earlier than smaller-capacity networks.Capacity is varied through widths W ∈ {8, 16, 32, 48, 64}, corresponding to p ∈ {0.26, 1, 4, 9, 16} × 10^6 parameters.

3 Training dynamics of a Random Features Network

The RFNN analysis links score-learning dynamics to the spectrum of a matrix U, revealing separated generalization and memorization timescales in the overparameterized regime. The first relaxation produces generalization, while slower modes drive overfitting and memorization.

  • Model and dynamics: The RFNN freezes Gaussian first-layer weights and learns second-layer weights by gradient descent in a high-dimensional limit.The model studies empirical score matching at fixed diffusion time, with training and test losses evaluated during learning.
  • Spectral structure: The spectrum of U contains well-separated bulks ρ1 and ρ2 when ψp ≫ ψn ≫ 1.Their supports scale respectively as ψp/ψn and ψp, while an additional delta peak is shown near λ = s² in the spectrum illustration.
  • Memorization timescale: τmem is associated with ρ1 and scales as ψp/(∆tλmin), becoming of order ψn/∆t when p ≫ n.During this later stage, training loss decreases while test loss increases and Lgen grows, indicating overfitting and memorization.
  • Capacity threshold: For n > n*(p) = p, asymptotic generalization loss is reduced and the model is no longer expressive enough to memorize the training data.Some overfitting remains, so this regime is described as architectural regularization rather than complete elimination of generalization loss.

4 Conclusions

The paper concludes that training dynamics provide implicit regularization against memorization in diffusion models. It supports this picture across realistic, synthetic, and analytically tractable settings while identifying optimizer, conditioning, capacity, and data-distribution boundaries.

  • Conclusions: Two well-separated timescales yield high-quality novel samples at τgen and memorization beyond τmem.The gap grows with training-set size, creating a broad early-stopping window for generalization.
  • Conclusions: The implicit dynamical regularization prevents memorization in highly overparameterized diffusion models.The authors state that the underlying score-learning mechanism is also common to stochastic interpolants and flow matching.
  • Evidence base: The results are demonstrated in realistic settings, controlled synthetic data, and analytically tractable models.The conclusion frames the evidence across these three types of settings rather than a single experimental setup.
  • Limitations and future work: Most practical diffusion models use Adam, whereas the derivation assumes SGD; optimizer-dependent shifts in the timescales remain to be studied.Additional experiments show that the two timescales still arise with Adam, but with fewer optimization steps.
  • Limitations and future work: The experiments are unconditional, and the absolute dependence of τmem and τgen on conditioning remains open.Toy conditional experiments verify that τmem still scales with n.
  • Limitations and future work: The numerical study covers p from 1M to 16M, leaving the full (n, p) phase diagram incompletely mapped.The authors identify wider capacity ranges as necessary for characterizing expressivity effects on dynamical regularization.
  • Limitations and future work: The theoretical analysis uses well-controlled data and score models, with richer distributions, structured architectures, and heavy-tailed data left for future work.These extensions are proposed to further characterize implicit dynamical regularization.
  • Practical implications: The authors propose early stopping and capacity control as practical guidelines for reducing memorization, especially in data-scarce domains.The conclusion notes that some industrial diffusion models still exhibit partial memorization.

Supplementary Material (SM)

The supplementary material supplies derivations and additional experiments supporting the main text. It covers numerical details, simplified models, formal theorem proofs, and the RFNN experiments.

  • Supplementary Material: The supplementary material provides detailed derivations and additional experiments supporting the main text.It is organized into sections addressing numerical experiments, simplified score and data models, proofs, and RFNN experiment details.
  • Supplementary Material: Section A gives details of the numerical experiments from Section 2.
  • Supplementary Material: Section B adds numerical experiments on simplified score and data models.
  • Supplementary Material: Section C provides formal proofs of the main theorems in Section 3.
  • Supplementary Material: Section D gives further details on the numerical experiments in Section 3.

A.1 Details on the numerical setup

The numerical study uses controlled CelebA experiments with U-Net DDPMs, evaluates generation and memorization across training time, and tests whether the transition persists under full-batch training and Adam.

  • CelebA images are center-cropped, converted to grayscale, and varied across training sets of n = 128 to 32768 without augmentation.
  • The score model is a three-level U-Net with channel multipliers {1, 2, 4}, attention at two coarsest resolutions, and base width W between 16 and 64.
  • Models use SGD with learning rate η = 0.01, momentum β = 0.95, and batch size B = min(n, 512), typically for at least 2M steps.
  • At τ = 190K, the W = 32 model has fmem = 0% and test FID 35.1; by τ = 1.62M, fmem reaches 47.2%.
  • Full-batch updates preserve τmem ∝ n, while Adam preserves the two-phase dynamics and shifts both τgen and τmem to fewer steps.With B = n, every sample is seen exactly τ times at fixed τ, so repetition does not explain the scaling.

B.1 Settings

The controlled GMM study evaluates a residual score network, analytical generalization metrics, and conditional generation to test the reported transition beyond natural-image experiments.

  • The data are two balanced d-dimensional Gaussians centered at ±µ with unit covariance, with µ = 1d enabling analytical computation of the generalization score.
  • Figure 7 compares FID and fmem over τ for varying n, using τ/n rescaling to assess memorization scaling under B = n and Adam.
  • The score model projects Rd to RW, applies three residual blocks with width W, skip connections, and layer normalization, then projects back to Rd.
  • Training uses full-batch SGD with η = 6 × 10^-3, momentum β = 0.95, n ∈ {128, 256, 512, 1024, 2048, 4096}, and up to 4M updates.
  • KL divergence measures agreement between the generated distribution Pθ and known data distribution P0 alongside memorization fraction fmem.
  • Conditional GMM experiments use classifier-free guidance and test whether rescaling training time by n collapses memorization curves across n ∈ {256, 512, 1024}.

C Proofs of the analytical results

The analytical appendix derives Gaussian-equivalent matrix descriptions and connects their spectra to the training dynamics in a high-dimensional random-features setting.

  • The replica-method derivation is not mathematically rigorous, but the final result can also be obtained rigorously through a free random matrix approach.
  • The training-dynamics timescales are determined by inverse eigenvalues of ΔtU/ψp, linking spectral analysis to temporal behavior.
  • The derivation assumes sub-Gaussian data tails and bounded covariance spectral norm.
  • The Gaussian Equivalence Principle replaces the spectrum of the original matrix with that of a Gaussian-equivalent construction in the high-dimensional limit.
  • The Gaussian-equivalent bulk spectrum has support of order ψp under the stated positive-definite covariance assumptions.

C.4 Proof of Theorem 3.1

The proof of Theorem 3.1 computes the spectrum of the relevant matrix through a replica partition function and saddle-point equations.

  • The proof begins by expressing the logarithm through the replica trick and defining a partition function for integer replica number s.
  • Dataset dependence is incorporated by changing variables Xν = Σ^1/2Zν, with Gaussian Zν having unit variance.
  • Order parameters Qab and analogous variables enforce replica constraints through delta-function representations.
  • In the high-dimensional limit, the partition function is evaluated at a saddle point, followed by differentiation with respect to the conjugate variables.
  • The derivation recovers the previously known isotropic-covariance result when Σ = Id.
  • The solution q* of the saddle-point equations corresponds to the Stieltjes transform of the eigenvalue density ρ.

C.5 Proof of Theorem 3.2

The proof decomposes the spectrum into bulk components and a delta peak, identifying which spectral structures contribute to the losses and how their scales separate.

  • Second bulk: The second bulk ρ2 matches the asymptotic eigenvalue density of the population covariance and is independent of ψn, with support on scale ψp.
  • Scale separation: For ψp ≫ ψn, the supports of ρ1 and ρ2 lie on scales ψp/ψn and ψp, respectively, so the two bulks are well separated.
  • Delta peak: When ψp ≥ ψn + 1, the spectrum contains a delta peak because n + d constraints admit nontrivial kernel vectors.
  • Delta peak: Eigenvectors associated with isolated eigenvalues do not contribute to train or test losses under zero initialization.The loss terms involve traces such as Tr(A^T A...) and Tr(AW), while the relevant eigenvectors lie in the null space of W.
  • First bulk: The proof identifies a first spectral bulk whose edges arise where the quadratic equation’s discriminant vanishes.This bulk is denoted ρ1.
  • First bulk: At leading order, the first bulk is independent of the population covariance spectrum ρΣ.

C.6 Dynamics on the fast timescales

The fast-timescale analysis shows that the second spectral bulk rapidly relaxes, making train and test losses nearly equal while generalization loss vanishes asymptotically.

  • Fast-timescale regime: On timescales 1 ≪ τ ≪ ψn, the train and test losses satisfy the fast-timescale asymptotic described by Proposition C.2.
  • Spectral dynamics: The spectrum has two relevant bulks: one with eigenvalues of order ψp/ψn and another with eigenvalues of order ψp.
  • Spectral dynamics: For the second bulk, the exponential relaxation factor is e^(-λ Δt τ), whereas the first-bulk factor remains approximately one in this regime.
  • Generalization loss: The generalization-loss terms from the second bulk cancel or vanish asymptotically, yielding Ltrain ≃ Ltest on the fast timescale.The cancellation uses shared generalized eigenvalue structure, while the remaining difference is bounded through the operator norm of U − Ũ.
  • Loss scale: The fast-timescale loss contribution remains a positive O(1) quantity because the relevant sum has d terms and eigenvalues of order ψp.

D Numerical experiments for Random Features

The numerical experiments test the random-features predictions across diffusion times, data variance, score error, spectral behavior, and optimization, including Adam.

  • Effect of diffusion time: Rescaling training time by τmem = ψp/(Δt λmin) collapses loss curves across diffusion times t in Figure 11.
  • Effect of diffusion time: In the overparameterized regime p ≫ n, the ratio τmem/τgen is independent of t, while p ∼ n produces nontrivial scaling.
  • Scope of interpretation: The experiments report increasingly separated timescales, but leave unresolved whether this behavior reflects specific properties of the learned score or proximity to the interpolation threshold.
  • Data variance: The predicted memorization timescale holds for data with variance σx ≠ 1, as shown by train and test loss curves in Figure 12.
  • Spectrum: Figure 14 compares theoretical and empirical spectra of U for ψp = 64, ψn ∈ {8, 32}, t = 0.01, and d = 100.
  • Optimization: For Adam optimization, random-features experiments retain the linear scaling of memorization time with n.
Loading 2505.17638v2…