Generalization of Diffusion Models Arises with a Balanced Representation Space
Zekai Zhang, Xiao Li, Xiang Li, Lianghe Shi, Meng Wu, Molei Tao, Qing Qu
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 dataset 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 our 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 modelling. Code is available at https://github.com/la0ka1/diffusion-gen-from-rep.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 2f237521-8d7e-4257-9151-3df08e817c23Cited by top-tier papers4
- Neural Network-Based Score Estimation in Diffusion Models: Optimization and GeneralizationYinbin Han, Meisam Razaviyayn, Renyuan XuICLR 2024 · 33 citations
- Multi-Subspace Multi-Modal Modeling for Diffusion Models: Estimation, Convergence and Mixture of ExpertsRuofeng Yang, Yongcan Li, Bo Jiang, Cheng Chen et al.ICLR 2026 · 4 citations
- Evaluating the Representation Space of Diffusion Models via Self-Supervised PrinciplesXiao Li, Yixuan Jia, Zekai Zhang, Xiang Li et al.ICML 2026
- Finding DoRI: Discovery of Retained Images in Diffusion ModelsAntoni Kowalczuk, Dominik Hintersdorf, Lukas Struppek, Kristian Kersting et al.ICML 2026
Builds on60
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Denoising Diffusion Implicit ModelsJiaming Song, Chenlin Meng, Stefano ErmonICLR 2021 · 11,743 citations
- Adding Conditional Control to Text-to-Image Diffusion ModelsLvmin Zhang, Anyi Rao, Maneesh AgrawalaICCV 2023 · 6,759 citations
- Scalable Diffusion Models with TransformersWilliam Peebles, Saining XieICCV 2023 · 5,568 citations
Related papers
- Deconstructing Denoising Diffusion Models for Self-Supervised LearningXinlei Chen, Zhuang Liu, Saining Xie, Kaiming HeICLR 2025 · 17 citations
- On the Edge of Memorization in Diffusion ModelsSam Buchanan, Druv Pai, Yi Ma, Valentin De BortoliNeurIPS 2025 · 25 citations
- Understanding Generalizability of Diffusion Models Requires Rethinking the Hidden Gaussian StructureXiang Li, Yixiang Dai, Qing QuNeurIPS 2024 · 45 citations
- Does Generation Require Memorization? Creative Diffusion Models using Ambient DiffusionKulin Shah, Alkis Kalavasis, Adam R. Klivans, Giannis DarasICML 2025
- Why Diffusion Models Don't Memorize: The Role of Implicit Dynamical Regularization in TrainingTony Bonnaire, Raphaël Urfin, Giulio Biroli, Marc MézardNeurIPS 2025 · 93 citations
