Generalization of Diffusion Models Arises with a Balanced Representation Space
This addresses the risk of overfitting in diffusion models for generative AI, offering insights to improve generalization, though it is incremental as it builds on existing representation learning concepts.
The paper tackles the problem of diffusion models memorizing training data instead of generalizing, by analyzing memorization and generalization through representation learning, and shows that generalization arises with balanced representations, validated on real-world models with practical applications like detection and editing.
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.