English

Generalization of Diffusion Models Arises with a Balanced Representation Space

Machine Learning 2026-02-12 v2 Computer Vision and Pattern Recognition

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.

Keywords

Cite

@article{arxiv.2512.20963,
  title  = {Generalization of Diffusion Models Arises with a Balanced Representation Space},
  author = {Zekai Zhang and Xiao Li and Xiang Li and Lianghe Shi and Meng Wu and Molei Tao and Qing Qu},
  journal= {arXiv preprint arXiv:2512.20963},
  year   = {2026}
}

Comments

Accepted at ICLR 2026. 40 pages, 19 figures. The first two authors contributed equally

R2 v1 2026-07-01T08:39:35.809Z