🤖 machine learning

Understanding diffusion models requires rethinking (again) generalization

This position paper argues that understanding generalization in diffusion models requires a new theoretical framework focused on the pre-memorization learning phase rather than explaining the absence of memorization, a stance supported by empirical findings on CIFAR-10 and a set of proposed open questions.

Original authors: Pierre Marion, Yu-Han Wu

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Pierre Marion, Yu-Han Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Why AI Art Doesn't Just Copy-Paste

Imagine you are teaching a student to paint. In traditional school (supervised learning), if the student memorizes the textbook perfectly, they might still be able to answer new questions on a test. But in the world of Diffusion Models (the AI behind tools like DALL-E or Stable Diffusion), the rules are different.

If this AI student memorizes the textbook (the training data) too perfectly, it stops being creative. Instead of painting a new picture of a cat, it just photocopies the exact cat from the textbook.

This paper argues that for a long time, researchers have been asking the wrong question. They've been asking, "Why doesn't the AI memorize the training data?" The authors say: "Stop asking that! We already know the answer."

Instead, we should be asking: "What is the AI actually learning before it starts memorizing?"

The Three Theories (The "Why" We Thought We Knew)

Before this paper, experts had three main guesses about why AI models don't just copy-paste their training data:

The "Small Backpack" Theory (Capacity): The model is too small to hold all the pictures in its memory, so it has to generalize.
The "Rough Ride" Theory (Optimization): The way the AI learns (the math behind the scenes) accidentally stops it from memorizing, kind of like how a bumpy road prevents a car from reaching a specific destination.
The "Specialized Brain" Theory (Architecture): The AI's internal structure is built in a way that naturally prefers patterns over exact copies.

The authors say these theories are all partially true, but they are missing the bigger picture.

The Experiment: A Controlled Art Class

To test these ideas, the researchers set up a "mini art class" using a small dataset (CIFAR-10, which is like a box of 10 types of simple toy images). They trained AI models of different sizes on different amounts of data, watching them closely step-by-step.

They looked for two things:

Novelty: Is the AI making new pictures, or just copying old ones?
Fidelity: Are the pictures good quality and realistic?

The Surprising Discoveries

Here is what they found that changed their minds:

1. The "Double Descent" Mystery
In traditional learning, as a model gets smarter, its errors go down, then up (if it memorizes), then down again. But in these AI models, the "quality" of the images went down, then up, then down again even before the model started memorizing.

Analogy: Imagine a student taking a test. Usually, they get better as they study. But here, the student's answers got worse in the middle of studying, then got great, then got worse again, all while they were still learning the concepts and not just memorizing the answers. The researchers don't know exactly why this happens yet.

2. The "Train-Test Gap" is a Lie
In normal AI, we check if a model is cheating by comparing how well it does on the "practice test" (training data) vs. the "real test" (new data).

The Finding: In diffusion models, the AI performs almost identically on both. The standard math tools used to detect cheating (memorization) don't work here. You can't tell if the AI is copying or creating just by looking at the numbers; you have to look at the actual images.

3. The "Time" Factor
They confirmed that memorization does happen, but it takes a very long time.

The Finding: If you have a huge dataset (like millions of photos), the AI would need to train for an impossibly long time before it starts copying. That's why real-world AI doesn't usually memorize. It simply runs out of time (or money) before it gets to the "copying" phase.

4. The "Tuning Knobs" (Batch Size & Learning Rate)
They tweaked the settings of how the AI learns.

Small Batches: When the AI learns in small groups (small batches), it makes better pictures during the learning process, even though it doesn't change when it starts memorizing.
Big Learning Rates: When the AI learns aggressively (high learning rate), it also makes better pictures during the process.
The Twist: These settings help the AI create better images while it is learning, but they don't necessarily stop it from memorizing later. It's like a student who studies very hard and gets great grades on the practice test, but if they study too long, they eventually just memorize the answer key.

The New Conclusion: What Should We Study?

The paper concludes that we have solved the mystery of "Why doesn't it memorize?" (Answer: It takes too long, and we stop training before that happens).

The new, hard question is: "What is happening in the 'Pre-Memorization' phase?"

The Analogy: Imagine a chef learning to cook. We used to worry about when they would start stealing recipes from the book. We now know they won't steal the recipes unless they cook for 100 years.
The Real Question: We need to understand what happens in the first 99 years. How do they learn to combine flavors to make a new dish? How do they learn the "soul" of the food rather than just the recipe?

Summary of Open Questions

The authors leave us with three big puzzles for the future:

Better Rulers: We need new ways to measure if an AI is "copying" or "creating," because our current math tools are blind to the difference.
The Learning Journey: We need to understand how the AI's "learning style" (the math settings) changes its creativity while it is learning, not just at the very end.
The Shape of Data: We need to understand how the shape of the data itself (the geometry of the pictures) helps the AI learn to be creative.

In short: Stop worrying about the AI stealing the textbook. Start studying how it learns to write its own story.

Technical Summary: Understanding Diffusion Models Requires Rethinking Generalization

Problem Statement

The paper argues that understanding generalization in diffusion models cannot rely on classical statistical learning theory (e.g., VC dimension, Rademacher complexity) or the "benign overfitting" paradigm developed for supervised learning. In supervised learning, interpolation (memorization) and generalization can coexist. In contrast, for diffusion models trained via denoising score matching, reaching the global minimum of the empirical risk corresponds to the score function of the noised empirical measure—a mixture of Gaussians centered at training points. Sampling from this minimum produces near-exact copies of training data rather than novel samples.

While practical diffusion models generate diverse, novel data, implying some form of regularization prevents full memorization, the theoretical explanations for this phenomenon remain fragmented. Existing literature attributes this to three distinct factors: (i) capacity limitations relative to dataset size, (ii) implicit regularization from optimization dynamics (e.g., early stopping, learning rate), and (iii) architectural inductive biases. The paper posits that the interactions between these factors are poorly understood and that the field has focused too heavily on why models do not memorize, rather than what they learn during the pre-memorization phase.

Methodology

To address these gaps, the authors conduct an empirical study using U-Net architectures trained on the CIFAR-10 dataset. The study is designed as an intermediate-scale investigation, bridging the gap between theoretical analyses on Gaussian data and large-scale practical systems.

Experimental Setup:

Data: The first $N$ samples of the CIFAR-10 training set are used, with $N$ varying across $\{2048, 4096, 8192, 16384\}$ . Evaluation uses a fixed test set of 5,120 samples.
Architecture: U-Nets with varying parameter counts ( $P \in \{3M, 8M, 26M, 93M, 201M\}$ ).
Training: Rectified flow formulation with 250 sampling steps and classifier-free guidance strength of 0.0.
Hyperparameter Sweeps: Systematic variation of dataset size ( $N$ ), model size ( $P$ ), batch size ( $B \in \{16, 32, 64, 128\}$ ), and learning rate ( $\eta \in \{10^{-5}, 5\cdot10^{-5}, 10^{-4}, 4\cdot10^{-4}\}$ ).

Metrics:
The authors track four key quantities on both train and test splits:

Denoising Score Matching Loss: Computed on a single batch (train) and full set (test).
Memorization Score: The ratio of generated samples with high cosine similarity (>0.6) to training data in a self-supervised feature space (SSCD).
Sliced Wasserstein Distance: Computed between generated samples and reference splits in Inception feature space.
FID (Fréchet Inception Distance): Computed between generated samples and reference splits.

The study defines memorization time ( $\tau_{mem}$ ) as the point where the test loss begins to diverge from the training loss. It distinguishes between novelty (absence of memorization, $\tau < \tau_{mem}$ ) and fidelity (quality/distributional accuracy).

Key Contributions and Results

1. Re-evaluating the "Why" of Non-Memorization

The authors confirm that the memorization transition is systematic and observable. They validate that the time to memorization ( $\tau_{mem}$ ) scales linearly with dataset size. This supports the conclusion that the question of why large-scale practical models do not memorize is largely resolved: the required training time to reach the memorization phase simply exceeds reasonable budgets. The authors argue the field should pivot from explaining the absence of memorization to analyzing the pre-memorization phase.

2. Limitations of Standard Metrics

A critical finding is that standard distributional metrics (Sliced Wasserstein, FID) fail to disentangle novelty from fidelity.

Train-Test Gap: Unlike supervised learning, the train-test gap in distributional distance is uninformative. Because the training and test sets are drawn from the same distribution, the distance from generated samples to either set remains nearly identical even as the model transitions from generating novel images to copying training data.
Double Descent: The authors observe a "double descent" phenomenon in distributional distances (FID and Sliced Wasserstein) for both train and test sets. This contradicts classical double descent theories where train and test metrics diverge. Instead, the model reaches an intermediate state with low training loss but poor distributional alignment, suggesting the phenomenon is driven purely by optimization dynamics rather than capacity arguments.

3. Distinct Roles of Dataset and Model Size

The study reveals that dataset size and model size play opposing roles in training dynamics:

Dataset Size: Increasing $N$ delays the onset of memorization and improves peak fidelity.
Model Size: Increasing $P$ accelerates the onset of memorization but also improves peak fidelity.
This suggests that generalization in diffusion models cannot be explained solely by the ratio of model capacity to sample size; the trajectory through parameter space shaped by optimization dynamics is a distinct factor.

4. Optimization Hyperparameters and Fidelity

The authors find that optimization hyperparameters significantly influence fidelity before memorization occurs:

Batch Size: Smaller batch sizes lead to better peak fidelity (lower minimum distributional distance) without significantly altering the memorization rate or onset time. This aligns with the "minimum stability" literature but manifests along the training trajectory rather than at convergence.
Learning Rate: Larger learning rates improve peak fidelity and delay the normalized memorization onset. This supports the "edge of stability" hypothesis but again, the effect is observed during the trajectory, not just at the final solution.

Significance and Position

The paper claims that the theoretical framework for generalization in diffusion models requires a fundamental shift. The authors argue that:

The "Why" is Solved: The lack of memorization in large-scale applications is explained by the linear scaling of memorization time with dataset size combined with early stopping.
The "What" is the Open Question: The primary theoretical challenge is understanding what is learned during the pre-memorization phase and how novelty and fidelity jointly emerge from the interplay of optimization, architecture, and data geometry.
New Framework Needed: Classical generalization gaps (in loss or distribution space) are insufficient for generative models. The training loss (score matching) is an indirect proxy for sample quality, and standard metrics fail to distinguish between a model that generalizes and one that memorizes.

The authors conclude that future research must focus on developing metrics that capture the copy-vs-generalize distinction and formulating a theory of implicit regularization that describes the properties of intermediate iterates in score matching, rather than just the final converged solution.