Generalization of Diffusion Models Arises with a Balanced Representation Space

TL;DR

This work tackles how diffusion models generalize without merely memorizing training data. By analyzing a nonlinear two-layer ReLU denoising autoencoder under a mixture-of-Gaussians data setup, the authors identify memorization as storage of exact training samples in weights (producing spiky activations) and generalization as learning local data statistics (producing balanced, semantic codes). They show three regimes—memorization, generalization, and a hybrid with imbalanced data—and validate that similar representation structures emerge in real-world diffusion models. Leveraging these insights, they propose a representation-based memorization detector and a training-free representation-space steering method for interpretable image editing, highlighting that learning good representations is central to meaningful, controllable generative modeling with practical implications for privacy and reliability.

Abstract

Diffusion models excel at generating high-quality, diverse samples, yet they risk memorizing training data when overfit to the training objective. We analyze the distinctions between memorization and generalization in diffusion models through the lens of representation learning. By investigating a two-layer ReLU denoising autoencoder (DAE), we prove that (i) memorization corresponds to the model storing raw training samples in the learned weights for encoding and decoding, yielding localized "spiky" representations, whereas (ii) generalization arises when the model captures local data statistics, producing "balanced" representations. Furthermore, we validate these theoretical findings on real-world unconditional and text-to-image diffusion models, demonstrating that the same representation structures emerge in deep generative models with significant practical implications. Building on these insights, we propose a representation-based method for detecting memorization and a training-free editing technique that allows precise control via representation steering. Together, our results highlight that learning good representations is central to novel and meaningful generative modeling.

Figures (19)

• Figure 1: Overview of theoretical findings. We analyze a two-layer ReLU DAE (middle row) to explain diffusion model behavior. We show that: (i) memorization occurs by storing raw samples in weights (top left), yielding spiky activations (bottom left); (ii) generalization arises by learning data statistics (top right), producing balanced, semantic codes (bottom right); and (iii) real-world models exhibit a hybrid regime due to data imbalance (top center). These insights unify representation and distribution learning, with empirical validation in \ref{['fig:teaser']}.
• Figure 2: Diffusion models generalize while learning benign internal representations. Activations from intermediate network layers form a representation space, within which distinct patterns emerge: memorized samples produce spiky representations that make them detectable, whereas novel generations yield balanced, information-rich representations that support controllable generation via representation steering.
• Figure 3: Sampling with Mem./Gen. ReLU DAEs.Left: sampling with a set of memorized ReLU DAE produces duplications of training images. Right: sampling with generalized DAEs produces novel images. Details for training and sampling are provided in \ref{['app:dae train/sampling']} and single-step denoising results are shown in \ref{['app:denoising']}
• Figure 4: Verification of \ref{['cor:dae-mem']} and \ref{['cor:dae-gen']}. We visualize the learned encoder matrix $\bm{W}_1$ of a ReLU DAE trained with noise level $\sigma=0.2$. When trained on 5 CelebA face images, the model stores training samples in its columns, matching \ref{['cor:dae-mem']}. When trained on 10,000 images, the model generalizes and captures data statistics, consistent with \ref{['cor:dae-gen']}. Empirically, the same behavior holds for larger noise, up to $\sigma=5$; additional results are in \ref{['app:thm verification']}.
• Figure 5: Mem./Gen. representations in ReLU DAEs.Top: Memorized vs. generalized samples can be separated by the standard deviation (Std) of their representations: memorized models produce spiky, high-Std features, whereas generalized models do not. Bottom: Representation of a single training data. The memorized model exhibits large outlier activations (high Std); the generalized model yields a more balanced representation (lower Std), consistent with our theory. All models use $\sigma=0.2$. Left: CelebA. Right: MoG. See \ref{['app:dae train/sampling']} for details.

Theorems & Definitions (18)

• Definition 3.1: $(\alpha,\beta)$-Separability of Training Data
• Theorem 3.1: Block-wise Structure of Local Minimizers in the DAE Loss
• Corollary 3.2: Memorization in Overparameterized DAEs
• Corollary 3.3: Generalization in Underparameterized DAEs
• Corollary 3.4: DAE memorizes duplicates and generalizes on well-sampled modes
• Lemma C.1: Global minimizers of Regularized LAE
• proof
• Definition C.1: $(\alpha,\beta)$-Separability of Training Data
• Theorem C.2: Restatement of \ref{['thm:dae']}
• proof