Technical Summary: Denoising Score Matching with Random Features

Problem Statement

Diffusion models have achieved remarkable success in generative tasks, yet the scientific community lacks a unified theoretical understanding of their generalization and memorization behaviors. Empirical observations suggest that these behaviors are influenced by model complexity, dataset size, and the number of noise samples used during training. Specifically, there is a concern regarding memorization, where a model's output closely resembles specific training data points, raising privacy and creativity concerns.

While prior theoretical works have analyzed diffusion models, they often rely on the empirical optimal score function (which inherently memorizes data) or assume regimes where the dataset size scales exponentially with dimension. There is a gap in understanding the proportional regime (where data size $n$, model complexity $p$, and dimension $d$ scale linearly) and the specific role of $m$, the number of noise samples per data point used in the Denoising Score Matching (DSM) loss function. This paper addresses the question: Can the phenomenon of memorization be theoretically demonstrated when using a parametric class of functions (specifically Random Features Neural Networks) for the score function?

Methodology

The authors analyze the learning process of a diffusion model under a simplified but analytically tractable theoretical setting:

1. Target Distribution: The target distribution $P_0$ is a $d$-dimensional Gaussian distribution, potentially supported on a lower-dimensional subspace (inspired by the manifold hypothesis).
2. Score Function Parameterization: The score function is modeled using Random Features Neural Networks (RFNNs). The first layer weights are fixed random matrices, while the second layer weights are learned.
3. Training Objective: The model minimizes the Denoising Score Matching (DSM) loss. Crucially, the analysis considers the parameter $m$, representing the number of noise samples averaged per data point in the loss function. The study focuses on two extremes:
   • $m = \infty$: The expectation over noise is computed exactly.
   • $m = 1$: A single noise sample is used per data point (empirical estimate).
4. Asymptotic Regime: The analysis operates in the high-dimensional limit where $d, n, p \to \infty$ while keeping the ratios $\psi_n = n/d$ and $\psi_p = p/d$ fixed.
5. Analytical Tools:
   • Gaussian Equivalence Principle (GEP): Used to replace the non-linear activation function outputs with equivalent Gaussian random variables, allowing the treatment of the feature matrix as Gaussian.
   • Linear Pencils Theory: A random matrix theory technique used to compute the traces of rational functions of random matrices, enabling the derivation of exact asymptotic learning curves (test and train errors).

Key Contributions

The paper provides the first theoretical study of the role of $m$ in memorization within diffusion models. The main contributions are:

1. Exact Asymptotic Expressions: The authors derive closed-form expressions for the test and train errors of the DSM minimizer in the high-dimensional limit for both $m=1$ and $m=\infty$.
2. Phase Diagram of Generalization vs. Memorization: By characterizing the errors, the authors identify distinct regimes based on $\psi_n$, $\psi_p$, and $m$:
   • A crossover transition occurs when $p = n$ (i.e., $\psi_p = \psi_n$).
   • Generalization dominates when the model is under-parameterized ($p < n$).
   • Memorization dominates when the model is over-parameterized ($p > n$).
3. The Role of $m$ (Noise Samples):
   • In the over-parameterized regime ($p > n$), increasing $m$ intensifies memorization. As $m$ increases, the learned score function converges closer to the empirical optimal score, causing the backward diffusion process to collapse onto training samples.
   • In the under-parameterized regime ($p < n$), increasing $m$ enhances generalization. The model's limited expressibility prevents it from overfitting to the noise, even with more noise samples.
4. Double Descent Phenomenon: The study confirms the presence of the double descent curve in the DSM setting (specifically for $m=1$), where test error peaks at the interpolation threshold ($\psi_p = \psi_n$).

Results

The theoretical findings are summarized in a phase diagram (Fig. 1 in the paper) and validated through numerical experiments:

• Theoretical Curves: The derived learning curves show that for $m=\infty$ and small $t$ (near the end of the reverse process), the train error decreases as $\psi_p$ increases, but the test error rises sharply when $\psi_p > \psi_n$, indicating memorization. Conversely, for $m=1$, the test error remains low even for large $\psi_p$, suggesting the model does not exhibit strong memorization behavior.
• Numerical Validation:
  • Synthetic Data: Experiments with Gaussian data and RFNN scores confirm that memorization increases as the ratio $n/p$ decreases and as $m$ increases in the over-parameterized regime.
  • Real Data: Experiments using U-Net architectures on Fashion-MNIST and MNIST datasets demonstrate that the theoretical predictions hold in practical settings. Memorization (measured by the retrieval of training samples during generation) increases as the number of training samples $n$ decreases relative to model parameters $p$, and increases with larger $m$ in over-parameterized models.

Significance and Claims

The paper claims to provide a first theoretical study of the role of $m$ in the memorization of diffusion models. The authors emphasize that their simplified setting (RFNNs, Gaussian data, independent estimators per time step) serves as a foundational step toward understanding the mechanisms of generalization and memorization.

Key claims regarding the paper's significance include:

• Mechanistic Insight: The work elucidates why memorization occurs: in over-parameterized regimes, minimizing the DSM loss with high $m$ forces the learned score to approximate the empirical optimal score, which drives the generative process toward training data points.
• Practical Relevance: The theoretical predictions align with empirical observations in practical models (U-Nets on real datasets), suggesting that the number of noise samples per data point is a critical hyperparameter controlling the trade-off between generalization and memorization.
• Theoretical Bridge: The study bridges the gap between empirical observations of memorization and rigorous statistical physics analysis, offering a precise framework to predict model behavior based on the ratios of data size, model complexity, and noise sampling frequency.

The authors remain modest, noting that their analysis relies on simplifications for tractability and that future work is needed to extend these results to more complex score models, non-Gaussian data, and explicit time-dependent score functions.