Two Calm Ends and the Wild Middle: A Geometric Picture of Memorization in Diffusion Models
Abstract
Diffusion models generate high-quality samples but can also memorize training data, raising serious privacy concerns. Understanding the mechanisms governing when memorization versus generalization occurs remains an active area of research. In particular, it is unclear where along the noise schedule memorization is induced, how data geometry influences it, and how phenomena at different noise scales interact. We introduce a geometric framework that partitions the noise schedule into three regimes based on the coverage properties of training data by Gaussian shells and the concentration behavior of the posterior, which we argue are two fundamental objects governing memorization and generalization in diffusion models. This perspective reveals that memorization risk is highly non-uniform across noise levels. We further identify a danger zone at medium noise levels where memorization is most pronounced. In contrast, both the small and large noise regimes resist memorization, but through fundamentally different mechanisms: small noise avoids memorization due to limited training coverage, while large noise exhibits low posterior concentration and admits a provably near linear Gaussian denoising behavior. For the medium noise regime, we identify geometric conditions through which we propose a geometry-informed targeted intervention that mitigates memorization.
Lay Summary
Modern AI-image generation methods often use something called a Diffusion Model, which is a process that starts with a noisy image and trains a neural network to incrementally remove the noise until only a realistic image remains. While powerful, a problem with Diffusion Models is that after training they may memorize, only returning images seen during training. This leads to potentially serious privacy and copyright concerns. In our paper we frame this problem using geometry, and classify model training into three regimes: small, medium, and large which depend on the amount of noise present in the images. We show that the behavior of training is calm in the small and large regimes but wild in the middle. We further identify a “danger zone” within the medium regime and show through experiments that training less heavily in this zone helps avoid memorization. Our work sheds new light on both the conditions governing memorization and the development of strategies used to mitigate it. This ultimately contributes to the goal of making AI-models safer and more reliable.