Generalization of Diffusion Models Arises from a Regularized Representation Space

Zekai Zhang, Xiang Li, Xiao Li, Molei Tao, Qing Qu

Keywords: Diffusion models; Memorization; Generalization; Representation learning; Autoencoders

TL;DR: Learning good representations is central to novel and meaningful image generation with diffusion models.

Abstract:

Diffusion models generate high-quality, diverse samples with great generalizability, yet when
overfit to the training objective, they may memorize training data. We analyze memorization
and generalization of diffusion models through the lens of representation learning. Using a
two-layer ReLU denoising autoencoder (DAE) parameterization, we prove that memorization
corresponds to the model learning the raw data matrix for encoding and decoding, yielding
spiky representations; in contrast, generalization arises when the model captures local data
statistics, producing balanced representations. We validate these insights by investigating
representation spaces in real-world unconditional and text-to-image diffusion models, where the
same distinctions emerge. Practically, we propose a representation-based memorization detection
method and a training-free editing method that allows precise control via representation steering.
Together, our results underscore that learning good representations is central to novel and meaningful
generation.

Submission Number: 75