Provable Separations between Memorization and Generalization in Diffusion Models

Abstract: Diffusion models have achieved remarkable success across diverse domains, but they remain vulnerable to memorization -- reproducing training data rather than generating novel outputs. This not only limits their creative potential but also raises concerns about privacy and safety. While empirical studies have explored mitigation strategies, theoretical understanding of memorization remains limited. We address this gap through developing a dual-separation result via two complementary perspectives: statistical estimation and network approximation. From the estimation side, we show that the ground-truth score function does not minimize the empirical denoising loss, creating a separation that drives memorization. From the approximation side, we prove that implementing the empirical score function requires network size to scale with sample size, spelling a separation compared to the more compact network representation of the ground-truth score function. Guided by these insights, we develop a pruning-based method that reduces memorization while maintaining generation quality in diffusion transformers.

• The paper demonstrates that the ground-truth score function does not minimize the empirical loss, establishing a non-vanishing Loss-Gap measured by the Fisher divergence.
• It shows that the empirical score function requires network capacity scaling with the sample size, while the ground-truth score admits a compact representation.
• The study validates that pruning and weight decay effectively mitigate memorization, enhancing generation quality and sample diversity.

Provable Separations between Memorization and Generalization in Diffusion Models

Overview and Motivation

This paper provides a rigorous theoretical and empirical analysis of memorization in diffusion models, establishing provable separations between memorization and generalization. The authors develop two complementary perspectives: (1) a statistical separation showing that the ground-truth score function does not minimize the empirical denoising loss, and (2) an architectural separation demonstrating that representing the empirical score function requires network size scaling with sample size, while the ground-truth score admits a compact representation. These results are validated through experiments and leveraged to propose a pruning-based mitigation strategy for diffusion transformers.

Statistical Separation: Loss Gap and Fisher Divergence

The central statistical result is that the empirical denoising score matching loss is minimized not by the ground-truth score function, but by the empirical score function derived from the training data. The authors formalize the Loss-Gap as the difference in denoising loss between the ground-truth and empirical score functions, and show that this gap is exactly the Fisher divergence between the empirical and true forward marginals. Notably, this gap does not vanish with polynomially many samples in high dimensions, especially for small diffusion times $t$ and large within-component variance.

For mixture models with well-separated components, the authors prove a lower bound on the expected Loss-Gap:

$$\mathbb{E}_{\mathcal{D}}[\text{Loss-Gap}] = \Omega\left(d \sigma_t^{-2} + \operatorname{tr}(\Sigma)\right)$$

for $t$ in a specified range and sample size $\log n = \mathcal{O}(d)$. This result implies that, in practical finite-sample regimes, strong optimizers and sufficiently expressive networks will bias training toward memorizing the empirical score function, especially in high dimensions and for small $t$.

Figure 1: Smaller $t$ leads to larger Loss-Gap; with insufficient sample size, the gap is non-negligible, confirming the statistical separation between memorization and generalization.

Architectural Separation: Network Complexity for Score Functions

The architectural analysis focuses on the representational requirements for neural networks to approximate the empirical and ground-truth score functions. The authors show that the empirical score function, which corresponds to a Gaussian mixture with $n$ components, requires network width and parameter count scaling linearly with $n$ for a fixed approximation error $\epsilon$. In contrast, the ground-truth score function for a sub-Gaussian Hölder density can be approximated with network size scaling as $\epsilon^{-d/(2\beta)}$, independent of $n$.

Formally, for ReLU networks $\mathcal{F}_1$ (empirical) and $\mathcal{F}_2$ (ground-truth), the required widths are:

• Empirical: $W_1 = \tilde{\mathcal{O}}(n \log^3 \epsilon^{-1})$
• Ground-truth: $$W_2 = \tilde{\mathcal{O}}(\epsilon^{-d/(2\beta)} \log^7 \epsilon^{-1})$$

This separation implies that, as $n$ increases, the empirical score function becomes increasingly difficult to represent, while the ground-truth score remains compact. The authors further analyze the Lipschitz continuity of the score functions, showing that the empirical score can have a large Lipschitz constant in the small-$t$ regime, making it highly irregular and harder to approximate with smooth networks.

Empirical Validation: Gaussian Mixtures and CIFAR-10

The theoretical findings are validated through extensive experiments on synthetic Gaussian mixture data and CIFAR-10. Key observations include:

• Increasing network size leads to a transition from underfitting, to generalization, and ultimately to memorization, where generated samples replicate training data.
• Larger networks and smaller training sets increase memorization ratio; higher data dimension reduces memorization due to the curse of dimensionality.
• Weight decay and network pruning are effective in mitigating memorization, especially when sample size is limited.

  Figure 2: Experimental results on Gaussian mixture data show the interplay between network size, sample size, and memorization ratio; proper weight decay and pruning reduce memorization and improve generalization.

On CIFAR-10, the authors propose a one-shot pruning method for diffusion transformers, targeting attention heads with low importance in the small-$t$ regime. This approach reduces memorization ratio while maintaining generation quality and diversity, as measured by FID, precision, and recall. The pruned models generate more diverse samples and are less likely to replicate training data.

Figure 1: Generated images from the original and pruned models; the pruned model exhibits greater diversity and less replication of training samples.

Implications and Future Directions

The results have significant implications for the design and deployment of diffusion models:

• Privacy and Copyright: The provable separation between memorization and generalization highlights the risk of training data leakage, especially in overparameterized models.
• Model Selection: Network size and regularization must be carefully chosen to balance generalization and avoid memorization, particularly in low-data or high-dimensional regimes.
• Mitigation Strategies: Pruning and weight decay are theoretically justified and empirically validated as effective defenses against memorization.

The theoretical framework is currently limited to sub-Gaussian mixture models and does not extend to heavy-tailed distributions. Scaling the pruning method to larger models and datasets remains an open challenge.

Conclusion

This paper establishes rigorous statistical and architectural separations between memorization and generalization in diffusion models, providing both theoretical bounds and empirical evidence. The findings clarify the mechanisms underlying memorization, quantify its persistence in practical regimes, and motivate principled mitigation strategies. Future work should extend these results to broader data distributions and further optimize pruning and regularization techniques for large-scale generative models.