Generalization of Diffusion Models Arises with a Balanced Representation Space

Code and figures for the ICLR 2026 paper: Generalization of Diffusion Models Arises with a Balanced Representation Space. Also a minimal repo for training and analyzing memorizing / generalizing diffusion models.

Repo Structure

• 1-Theory/ (CPU): Train and analyze two-layer ReLU denoising autoencoder and diffusion models.
• 2-Application/ (GPU recommended): Visualize and steer SD v1.4 representations.
• Figs/: plots generated by the notebooks and used in this preview.

Quickstart

pip install numpy torch torchvision diffusers transformers accelerate datasets scikit-learn pillow tqdm matplotlib seaborn

1-Theory (CPU)

Notebooks: ReLU_DAE.ipynb, ReLU_Diffusion.ipynb

Train ReLU-DAE on CelebA: visualize representations and weights for memorization vs. generalization (Figures 4-5).

A diffusion extension in the same toy setup, with time embeddings and sampling.

2-Application (GPU recommended)

Notebooks: SD_compare_reps.ipynb, SD_steering.ipynb

Stable Diffusion v1.4 + LAION representation structure and separation (Figures 6a/6b).

Representation steering: effective for generalized samples, ineffective for memorized samples (Figure 8).

Additional figure

Steering trajectory in representation space, showing separation between concepts/styles and how steering transfers across them.

Citation

@inproceedings{zhang2026balanceddiffusion,
  title={Generalization of Diffusion Models Arises with a Balanced Representation Space},
  author={Zhang, Zekai and Li, Xiao and Li, Xiang and Shi, Lianghe and Wu, Meng and Tao, Molei and Qu, Qing},
  booktitle={ICLR},
  year={2026}
}