When Diffusion Models Memorize: Inductive Biases in Probability Flow of Minimum-Norm Shallow Neural Nets (2506.19031v1)

Abstract: While diffusion models generate high-quality images via probability flow, the theoretical understanding of this process remains incomplete. A key question is when probability flow converges to training samples or more general points on the data manifold. We analyze this by studying the probability flow of shallow ReLU neural network denoisers trained with minimal $\ell^2$ norm. For intuition, we introduce a simpler score flow and show that for orthogonal datasets, both flows follow similar trajectories, converging to a training point or a sum of training points. However, early stopping by the diffusion time scheduler allows probability flow to reach more general manifold points. This reflects the tendency of diffusion models to both memorize training samples and generate novel points that combine aspects of multiple samples, motivating our study of such behavior in simplified settings. We extend these results to obtuse simplex data and, through simulations in the orthogonal case, confirm that probability flow converges to a training point, a sum of training points, or a manifold point. Moreover, memorization decreases when the number of training samples grows, as fewer samples accumulate near training points.

Inductive Biases and Memorization in Probability Flow of Minimum-Norm Shallow Neural Nets

This paper provides a rigorous theoretical and empirical analysis of the probability flow ODE in diffusion models, focusing on shallow ReLU neural network denoisers trained with minimal $\ell^2$ norm. The work addresses a central question in generative modeling: under what conditions does the deterministic sampling process in diffusion models converge to training samples (memorization) versus more general points on the data manifold (generalization)?

Theoretical Framework

The analysis is grounded in the paper of shallow ReLU denoisers that interpolate noisy samples with minimal representation cost. The authors consider the low-noise regime, where noisy samples are well-clustered around clean data points. The denoiser is trained to map each noisy sample exactly to its corresponding clean point, resulting in zero empirical loss. The main theoretical contributions are:

• Score Flow vs. Probability Flow: The paper distinguishes between the score flow ODE (fixed noise level) and the probability flow ODE (with a time-dependent noise schedule). For orthogonal datasets, both flows follow similar trajectories, but the probability flow can terminate at more general manifold points due to "early stopping" induced by the diffusion time scheduler.
• Characterization of Convergence Points: For orthogonal data, the set of stable stationary points of the score flow consists of all possible sums of training points (vertices of a hyperbox). The probability flow can converge to these vertices or to points on the boundary of the hyperbox, depending on the initialization and the diffusion time schedule.
• Extension to Non-Orthogonal Data: The results are generalized to datasets forming obtuse simplices and equilateral triangles, with analogous geometric interpretations for the set of possible convergence points.

Empirical Validation

The theoretical findings are validated through extensive simulations:

• Existence of Virtual Points: Empirical results confirm that the score flow and probability flow can converge to "virtual points"—sums of training points not present in the original dataset. The stability of these points decreases as the number of combined training points increases, but their absolute number grows combinatorially.
• Effect of Training Set Size: As the number of training samples increases, the proportion of generated samples converging to virtual points or manifold boundaries increases, indicating improved generalization and reduced memorization.
• Training Protocols: The minimum-norm constraint is enforced via an augmented Lagrangian method. Standard training with weight decay regularization approximates the minimum-norm solution, while omitting regularization leads to increased memorization.
• Robustness to Data Augmentation: When the training set is augmented with random points on the hyperbox boundary, the probability flow continues to converge predominantly to the boundary, supporting the geometric interpretation.

Numerical Results

Key empirical findings include:

• For orthogonal datasets in $d=30$ with $N=30$ training points, over 98% of pairwise virtual points are stable under score flow.
• In probability flow, a significant fraction of samples converge to virtual points or hyperbox boundaries, with the proportion increasing as $N$ grows.
• The use of dropout during training increases the MSE loss and leads to more samples being generated outside the hyperbox boundary.

Implications

Practical Implications

• Understanding Memorization in Diffusion Models: The results provide a precise geometric and dynamical characterization of when diffusion models memorize training data versus when they generalize. This has direct implications for privacy, as memorization can lead to training data leakage.
• Design of Denoisers: The analysis suggests that enforcing minimum-norm solutions (e.g., via weight decay) can bias the generative process toward generalization, producing novel combinations of training data rather than exact replicas.
• Sampling Strategies: The role of the diffusion time scheduler is critical; "early stopping" can be used to control the trade-off between memorization and generalization in practical sampling algorithms.

Theoretical Implications

• Geometry of Learned Manifolds: The convergence to hyperbox boundaries and virtual points provides a concrete geometric framework for understanding the implicit data manifolds learned by shallow denoisers.
• Extension to Deep Networks: While the analysis is restricted to shallow networks, the geometric phenomena observed may extend to deeper architectures, especially in high-dimensional settings where data points become approximately orthogonal.
• Connection to Associative Memory: The results align with prior work on overparameterized autoencoders and associative memory, further bridging the gap between generative modeling and memory retrieval dynamics.

Limitations and Future Directions

• Model Depth and Realism: The analysis is limited to shallow networks and assumes perfect interpolation in high-dimensional balls around each training point. Real-world diffusion models are deep and may not satisfy these assumptions.
• Variance-Preserving Processes: The current analysis focuses on variance-exploding (VE) processes. Extending the results to variance-preserving (VP) processes, which are more common in practice, is an important direction.
• High-Dimensional Manifolds: The assumption of interpolation over $d$-dimensional balls may not hold when the number of noisy samples is less than the input dimension. Future work could consider interpolation over lower-dimensional discs, reflecting the concentration of Gaussian noise in high dimensions.

Speculation on Future Developments

• Generalization in Deep Diffusion Models: As diffusion models scale, understanding and controlling the balance between memorization and generalization will become increasingly important, both for privacy and for creative generation.
• Geometric Regularization: Insights from this work may inform new regularization strategies that explicitly encourage generative models to sample from manifold boundaries or virtual points, enhancing diversity and reducing overfitting.
• Analytical Tools for Deep Models: The geometric and dynamical tools developed here could be adapted to analyze deeper architectures, potentially leading to new theoretical guarantees for generalization in deep generative models.

Summary Table: Convergence Behavior

Data Geometry	Score Flow Convergence	Probability Flow Convergence	Effect of Early Stopping
Orthogonal	Training points or sums (vertices of hyperbox)	Vertices or boundary points of hyperbox	Can halt at boundary
Obtuse Simplex	Subset of partial sums	Chord endpoints or points on chord	Can halt on chord
Equilateral Triangle	Face then vertex	Analogous, with possible early halt	Can halt on face

Conclusion

This work provides a comprehensive theoretical and empirical account of the inductive biases in the probability flow of minimum-norm shallow neural nets used in diffusion models. By elucidating the geometric structure of convergence points and the role of the diffusion time scheduler, the paper advances our understanding of memorization and generalization in generative modeling, with significant implications for both theory and practice.