Generalization of Diffusion Models Arises with a Balanced Representation Space

TL;DR

This paper explores how diffusion models generalize versus memorize by analyzing their representation structures, revealing that balanced representations lead to better generalization and proposing methods for detection and control.

Contribution

It provides a theoretical analysis linking representation structures to memorization and generalization in diffusion models, and introduces practical techniques for detection and editing.

Findings

Memorization involves storing raw training data in weights, leading to localized representations.

Generalization involves capturing local data statistics, resulting in balanced representations.

Representation structures observed in real-world diffusion models align with theoretical predictions.

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.