InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models

By Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, Volodymyr Kuleshov

While diffusion models excel at generating high-quality samples, their latent variables typically lack semantic meaning and are not suitable for representation learning. InfoDiffusion augments a diffusion model with a low-dimensional latent variable z that captures high-level factors of variation in the data. An encoder infers z from a clean image, and the denoiser is conditioned on z at every step. The learning objective is regularized with the mutual information between the image and z, which keeps the expressive diffusion decoder from ignoring z, and with a divergence between the aggregate posterior and a prior p(z) that can be continuous or discrete. InfoDiffusion learns disentangled, human-interpretable latent representations that are competitive with state-of-the-art generative and contrastive methods, while retaining the sample quality of diffusion models.

In this repo, we provide:

• InfoDiffusion
  1. A rectified-flow denoiser (velocity prediction, logit-normal time steps, 50-step Heun sampler, ODE inversion) with an ADM U-Net whose every residual block is modulated by z, and a half-U-Net encoder.
  2. The regularized objective with coefficients (λ, ζ): a per-sample KL term with weight 1 − ζ and an IMQ-kernel MMD with weight λ + ζ − 1. The main setting is (λ, ζ) = (0.1, 0.9).
  3. Latent dropout with classifier-free guidance on z, and a learned flow-matching prior over z.
  4. Discrete latents with finite scalar quantization (FSQ), plus the binary, relaxed-Bernoulli and LFQ variants of the paper.
• Evaluation: latent-quality probes, TAD, DCI, clean-FID with z from p(z) or the learned prior, and reconstructions.
• Reproduction: one registry entry per run of the paper (117 runs), the archived evaluations of all of them, the scripts that rebuild the paper's tables, and the DiffAE, AE, VAE, β-VAE and InfoVAE baselines.
• Checkpoints: 30 models (the main setting on five datasets with three seeds each, and the FSQ models) and 12 latent priors on Hugging Face.

Main Results

InfoDiffusion in the main setting, (λ, ζ) = (0.1, 0.9), against autoencoder and diffusion baselines (Tables 2 and 3 of the paper). InfoDiffusion entries are means ± standard deviations over three training seeds.

Latent quality (↑): 5-fold accuracy of a logistic regression on z (FFHQ: mean accuracy over age, gender and eyeglasses; CelebA: mean AUROC over 40 attributes), with d_z = 32.

Method	FashionMNIST	CIFAR10	FFHQ	CelebA
AE	0.819 ± 0.003	0.336 ± 0.005	0.532 ± 0.002	0.759 ± 0.003
VAE	0.796 ± 0.002	0.342 ± 0.004	0.488 ± 0.002	0.770 ± 0.002
β-VAE	0.779 ± 0.004	0.253 ± 0.003	0.510 ± 0.001	0.699 ± 0.001
InfoVAE	0.807 ± 0.003	0.357 ± 0.005	0.545 ± 0.001	0.757 ± 0.003
DiffAE	0.822 ± 0.006	0.442 ± 0.004	0.565 ± 0.001	0.817 ± 0.001
InfoDiffusion	0.856 ± 0.004	0.461 ± 0.003	0.628 ± 0.002	0.844 ± 0.005

Sample quality: FID (↓) of 10,000 samples against the training set, with d_z = 256. InfoDiffusion samples z from p(z) = N(0, I) or from the learned latent prior, with guidance weight w on z.

Method	FashionMNIST	CIFAR10	FFHQ	CelebA
AE	62.9 ± 2.1	169.4 ± 2.4	316.7 ± 0.4	90.4 ± 1.8
VAE	63.4 ± 1.6	177.2 ± 3.2	307.7 ± 0.1	94.3 ± 2.8
β-VAE	66.9 ± 1.8	183.3 ± 3.1	222.6 ± 0.2	99.8 ± 2.4
InfoVAE	55.0 ± 1.7	160.7 ± 2.5	195.2 ± 0.5	77.8 ± 1.6
DiffAE	5.6 ± 0.1	34.1 ± 0.1	12.1 ± 0.1	8.4 ± 0.1
InfoDiffusion, z ~ p(z)	22.5 ± 0.6	36.2 ± 1.1	19.8 ± 0.4	12.1 ± 0.2
w = 2	21.9 ± 0.7	33.7 ± 1.0	15.7 ± 0.7	11.1 ± 0.3
InfoDiffusion, learned latent	8.1 ± 0.4	25.7 ± 0.4	13.1 ± 0.5	7.6 ± 0.2
w = 2	6.8 ± 0.7	21.3 ± 0.7	9.8 ± 0.3	6.3 ± 0.2

Disentanglement (↑): DCI on 3DShapes (d_z = 10), and TAD with the number of captured attributes on CelebA (d_z = 32). "Image Qual." marks whether 16 samples look correct on 3DShapes (InfoDiffusion: with z from p(z) / from the learned latent prior).

Method	3DShapes DCI	3DShapes Image Qual.	CelebA TAD	CelebA Attrs
AE	0.219 ± 0.001	✗	0.042 ± 0.004	1.0 ± 0.0
VAE	0.276 ± 0.001	✗	0.000 ± 0.000	0.0 ± 0.0
β-VAE	0.281 ± 0.001	✗	0.088 ± 0.051	1.6 ± 0.8
InfoVAE	0.134 ± 0.001	✗	0.000 ± 0.000	0.0 ± 0.0
DiffAE	0.177 ± 0.001	✓	0.114 ± 0.033	2.8 ± 0.4
InfoDiffusion	0.296 ± 0.057	✗ / ✓	0.439 ± 0.134	5.7 ± 0.6

DiffAE is re-trained with its official code on our splits, and the FFHQ AE, VAE, β-VAE and InfoVAE with the 2023 InfoDiffusion code (baselines); for these, ± is the variation across cross-validation folds, evaluation subsamples or sample sets. The other baseline entries come from earlier runs with a separate training pipeline. With FSQ latents (Table 10), InfoDiffusion reaches a CelebA latent quality of 0.849 ± 0.003, TAD 0.370 ± 0.240 and 5.6 ± 0.9 captured attributes. Per-run metrics of all 117 runs are in results/runs/.

Code Organization

1. train.py: Training (single GPU or torchrun), with per-dataset presets (--preset) and resumable time budgets
2. evaluate.py: Latent encoding, probes, TAD and DCI, the learned latent prior, clean-FID and reconstructions
3. sample.py: Samples and reconstructions from a trained model
4. infodiff2/: The model and objective (model.py), U-Net and encoder (nets.py), rectified flow and samplers (flow.py), MMD and KL (losses.py), metrics (metrics.py), data (data.py)
5. configs/: One preset per dataset and model type; experiments.json registers every run of the paper
6. reproduce.py: Lists and launches the runs of the paper
7. scripts/: Data preparation, checkpoint download, the paper's tables, extra FID sample sets and a Slurm example
8. baselines/: The DiffAE baseline driver and the patch and runner for the 2023 code
9. results/: Archived evaluations behind the paper's tables, baseline numbers and the checkpoint manifest
10. tests/: CPU tests, including a short training and evaluation run
11. docs/: Data preparation, the reproduction guide and the checkpoint list

Getting Started

Create an environment with the tested dependencies (Python 3.10, PyTorch 2.5.1, CUDA 12.4) and run the CPU tests:

conda create --name infodiffusion python=3.10
conda activate infodiffusion
pip install -r requirements.txt
python -m unittest discover -s tests

Prepare the datasets as uint8 arrays under data/ (each one is downloaded from its public source; see docs/DATA.md for sources, splits and labels):

python scripts/prep_data.py fmnist --root data
python scripts/prep_data.py cifar10 --root data
python scripts/prep_data.py celeba --root data   # Hugging Face Datasets, several CPU-hours
python scripts/prep_data.py ffhq --root data     # thumbnails and ffhq-features-dataset labels
python scripts/prep_data.py 3dshapes --root data

Every model trains on one GPU. On an A100, the CIFAR10 generative model trains at about 16 iterations per second.

Checkpoints

We release the EMA weights of all models of the main setting (three seeds each), the learned latent priors of the generative models, and the FSQ models, on Hugging Face 🤗: yingheng/InfoDiffusion-checkpoints (5.9 GB). docs/CHECKPOINTS.md lists every file with its size, SHA256 and archived metrics.

Dataset	Latent model (d_z = 32)	Generative model (d_z = 256) + latent prior
FashionMNIST	pg_fm_rep_l01z09	pg_fm_gen_l01z09
CIFAR10	pg_c10_rep_l01z09	pg_c10_gen_l01z09
FFHQ	pg_ff_rep_l01z09	pg_ff_gen_l01z09
CelebA	pg_ca_rep_l01z09; FSQ: pg_ca_rep_fsq	pg_ca_gen_l01z09
3DShapes	pg_sh_rep_l01z09 (d_z = 10)	—

Seeds 1 and 2 add the suffixes _s1 and _s2.

# all files (verifies SHA256); --runs NAME ... downloads a subset
python scripts/download_checkpoints.py --output checkpoints
# samples from the learned latent prior with guidance w = 2
python sample.py --ckpt checkpoints/checkpoints/pg_ca_gen_l01z09.pt \
  --latent_prior checkpoints/latent_priors/pg_ca_gen_l01z09.pt --cfg 2 --out samples/celeba.png
# the paper's metrics for a released model
python evaluate.py --ckpt checkpoints/checkpoints/pg_ca_rep_l01z09.pt --data_root data \
  --out eval/pg_ca_rep_l01z09 --tasks latents,probe
python evaluate.py --ckpt checkpoints/checkpoints/pg_ca_gen_l01z09.pt --data_root data \
  --out eval/pg_ca_gen_l01z09 --tasks latents,fid --fid_modes prior,prior@2,learned,learned@2 \
  --latent_prior checkpoints/latent_priors/pg_ca_gen_l01z09.pt

Reproducing Experiments

Below, we describe how to reproduce the experiments in the paper. The main entry point is reproduce.py: each run of the paper is an entry of configs/experiments.json, and reproduce.py turns it into train.py and evaluate.py commands. Pass --dry-run to print them.

Training

The presets in configs/ are the paper's main setting, (λ, ζ) = (0.1, 0.9). Each dataset has a latent model, which the paper uses for latent quality and disentanglement (d_z = 32, d_z = 10 for 3DShapes; no latent dropout), and a generative model, which it uses for FID (*_gen, d_z = 256; latent dropout 0.1 for guidance on z):

python train.py --preset cifar10 --data_root data --out runs/pg_c10_rep_l01z09 --seed 0
python train.py --preset cifar10_gen --data_root data --out runs/pg_c10_gen_l01z09 --seed 0

Any flag overrides the preset, e.g. --latent fsq --stochastic 0 --w_kl 0 for FSQ latents or --stochastic 0 --w_kl 0 --w_mmd 0.01 for (λ, ζ) = (0.01, 1); configs/experiments.json has the exact flags of every run. Multi-GPU training uses torchrun --nproc_per_node=N train.py ... with the same global batch size.

Presets	Steps (epochs)	Width (dec. / enc.)	Pooling	Crops	t
fmnist, fmnist_gen	23,400 (50)	40 / 48	flatten	✓	logit-normal
cifar10, cifar10_gen	19,500 (50)	80 / 96	mean	✓	logit-normal
ffhq, ffhq_gen	24,600 (50)	80 / 96	flatten	✓	logit-normal
celeba, celeba_gen	63,600 (50)	80 / 96	mean / flatten	✓	logit-normal
3dshapes	16,875 (5)	40 / 48	flatten	✗	uniform

Pooling is applied to the encoder's last feature map before the linear head (CelebA: mean for d_z = 32, flatten for d_z = 256); crops are random resized crops of the encoder input; t is the distribution of training time steps.

To train and evaluate all 27 runs behind Tables 2 and 3:

python reproduce.py list --group main
# train, evaluate, and mark done (runs/NAME/eval.done); skips finished steps
python reproduce.py run pg_c10_rep_l01z09 pg_c10_rep_l01z09_s1 pg_c10_rep_l01z09_s2

scripts/slurm/pack.sbatch runs several experiments on one GPU and resubmits itself until they finish.

Evaluation

# latent model: probes (and TAD on CelebA, DCI on 3DShapes) and reconstructions
python evaluate.py --ckpt runs/pg_c10_rep_l01z09/ema_0019500.pt --data_root data \
  --tasks latents,probe,recon
# generative model: learned latent prior, then FID with z from p(z) and from the prior,
# without and with guidance (w = 2)
python evaluate.py --ckpt runs/pg_c10_gen_l01z09/ema_0019500.pt --data_root data \
  --tasks latents,probe,prior,fid --fid_modes prior,prior@2,learned,learned@2

Results are written to runs/<run>/eval/results.json. In the paper's tables, latent quality is probe.probe_paper_protocol[0], TAD and captured attributes are entries 0 and 2 of probe.tad_paper_protocol, DCI is probe.dci_seeds[0], and FID is fid.<mode>. The first FID computation for a dataset writes the training images as PNGs and computes the clean-FID reference statistics.

Tables

# mean ± std over seeds -> runs/paper_tables.json
python scripts/paper_tables.py runs
# InfoDiffusion rows of Tables 2-5 and 8-10
python scripts/paper_latex.py runs/paper_tables.json
# Tables 11 and 12, with the baselines
python scripts/paper_apph.py runs/paper_tables.json results/baselines.json results/iq_3dshapes.json
# the same tables from the archived evaluations
python scripts/paper_tables.py results/runs --archived --out paper_tables.json

Baselines and Other Experiments

The DiffAE baseline uses the official implementation, and the FFHQ AE, VAE, β-VAE and InfoVAE use the 2023 InfoDiffusion code; see baselines/README.md. The reproduction guide covers conditioning on z at the bottleneck only (Tables 8–9), discrete latents (Table 10) and the regularization coefficients (Tables 11–12), and explains the run names and training details.

Code History

This implementation was written for the revised version of the paper. The original 2023 implementation (DDPM-based) is at commit 75c7b31.

Acknowledgements

The U-Net follows ADM, FID uses clean-FID, and the DiffAE baseline runs the official DiffAE code. The TAD metric follows NashAE and DCI follows disentanglement_lib.

Citation

@inproceedings{wang2023infodiffusion,
  title     = {InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models},
  author    = {Wang, Yingheng and Schiff, Yair and Gokaslan, Aaron and Pan, Weishen and Wang, Fei and De Sa, Christopher and Kuleshov, Volodymyr},
  booktitle = {Proceedings of the 40th International Conference on Machine Learning},
  pages     = {36336--36354},
  year      = {2023},
  series    = {Proceedings of Machine Learning Research},
  volume    = {202},
  publisher = {PMLR}
}

License

The code is released under the MIT License; see LICENSE and THIRD_PARTY_NOTICES.md. The datasets are not covered and must be obtained under their own terms; CelebA and FFHQ are restricted to non-commercial research, and the released checkpoints, which were trained on them, inherit these restrictions.