Smoothing the Score Function for Generalization in Diffusion Models: An Optimization-based Explanation Framework

TL;DR

This work explains memorization in diffusion models through an optimization-driven lens: the empirical score is a sharp, softmax-weighted mixture of Gaussian components whose dominance tends to lock sampling to training points. Neural networks mitigate memorization by implicitly smoothing these weights and enabling local-manifold exploration. The authors introduce Noise Unconditioning, which unifies noise scales and reinterprets sampling as gradient ascent on a fixed Gaussian mixture, and Temperature Smoothing, which explicitly regulates weight smoothness via a temperature parameter. Through theoretical analysis and extensive experiments, they show that smoothing reduces memorization, accelerates sampling, and preserves generation quality, with strong evidence from Jacobian-spectrum analyses and ablations on CIFAR-10 and Cat Caracal. The framework links diffusion theory with optimization and geometry, offering practical routes to improve generalization in diffusion-based generation and suggesting extensions to latent diffusion and consistency models.

Abstract

Diffusion models achieve remarkable generation quality, yet face a fundamental challenge known as memorization, where generated samples can replicate training samples exactly. We develop a theoretical framework to explain this phenomenon by showing that the empirical score function (the score function corresponding to the empirical distribution) is a weighted sum of the score functions of Gaussian distributions, in which the weights are sharp softmax functions. This structure causes individual training samples to dominate the score function, resulting in sampling collapse. In practice, approximating the empirical score function with a neural network can partially alleviate this issue and improve generalization. Our theoretical framework explains why: In training, the neural network learns a smoother approximation of the weighted sum, allowing the sampling process to be influenced by local manifolds rather than single points. Leveraging this insight, we propose two novel methods to further enhance generalization: (1) Noise Unconditioning enables each training sample to adaptively determine its score function weight to increase the effect of more training samples, thereby preventing single-point dominance and mitigating collapse. (2) Temperature Smoothing introduces an explicit parameter to control the smoothness. By increasing the temperature in the softmax weights, we naturally reduce the dominance of any single training sample and mitigate memorization. Experiments across multiple datasets validate our theoretical analysis and demonstrate the effectiveness of the proposed methods in improving generalization while maintaining high generation quality.

Figures (33)

• Figure 1: Illustration of sampling behaviors with different score functions. (a) Standard diffusion empirical score function leads to memorization, where trajectories collapse directly to training points. (b) The neural network-learned score function enables generalization. (Interpolation experiment) (c) The noise unconditioning allows smoother score function weights and delayed collapse. (d) The temperature-based score function promotes manifold exploration and prevents premature collapse to training points.
• Figure 2: Illustration for the dataset of cat (left) and caracal (right). Note that there are no yellow or long-eared pet cats in the dataset.
• Figure 3: Qualitative comparison with different methods. (a, b, c) illustrates the generated images (left of red line) and their top-3 nearest neighbors of training samples in pixel space. (d) shows the generated images (2, 3, 4) and their corresponding closest training sample (1). (closest in both pixel and feature spaces)
• Figure 4: Sampling expansiveness ratio $\gamma_{\text{ex}}$ at different noise levels. Empirical conditioning scores are highly expansive at low noise levels, while unconditioning, temperature smoothing, and NN approximation all yield much smaller $\gamma_{\text{ex}}$.
• Figure 5: Generated images (top to bottom) for variants of SDE 1K NFE on datasets CIFAR-10 and CelebA. (pixel space KNN)