yingheng/InfoDiffusion-checkpoints

InfoDiffusion checkpoints

Companion weights for InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models (ICML 2023; arXiv:2306.08757). Code: github.com/yingheng-wang/InfoDiffusion. Project page: yingheng-wang.com/infodiffusion-page.

InfoDiffusion conditions a rectified-flow diffusion model on a low-dimensional latent z inferred by an encoder, and regularizes the objective with the mutual information between the image and z and a prior on z. All models use the paper's main setting, (λ, ζ) = (0.1, 0.9), and were trained with three seeds:

Dataset	Latent model (d_z = 32)	Paper metric (mean ± std over 3 seeds)	Generative model (d_z = 256)	FID, learned latent, w = 2
FashionMNIST	pg_fm_rep_l01z09	latent quality 0.856 ± 0.004	pg_fm_gen_l01z09	6.8 ± 0.7
CIFAR10	pg_c10_rep_l01z09	latent quality 0.461 ± 0.003	pg_c10_gen_l01z09	21.3 ± 0.7
FFHQ	pg_ff_rep_l01z09	latent quality 0.628 ± 0.002	pg_ff_gen_l01z09	9.8 ± 0.3
CelebA	pg_ca_rep_l01z09	AUROC 0.844 ± 0.005, TAD 0.439 ± 0.134	pg_ca_gen_l01z09	6.3 ± 0.2
CelebA, FSQ latents	pg_ca_rep_fsq	AUROC 0.849 ± 0.003, TAD 0.370 ± 0.240	—	—
3DShapes (d_z = 10)	pg_sh_rep_l01z09	DCI 0.296 ± 0.057	—	—

Seeds 1 and 2 carry the suffixes _s1 and _s2. Latent quality is the 5-fold accuracy (CelebA: mean AUROC over 40 attributes; FFHQ: mean accuracy over age, gender and eyeglasses) of a logistic regression on z. FID is clean-FID between 10,000 samples and the training set, with z sampled from the learned latent prior and guidance weight w = 2 on z; Tables 2 and 3 of the paper also give FIDs with z ~ N(0, I) and without guidance. These are archived evaluations of these exact files, not packaging-time reruns.

Contents and use

• checkpoints/<run>.pt (30 files): EMA weights, training step and training arguments. Files load with torch.load(..., weights_only=True); optimizer states and non-EMA weights are omitted.
• latent_priors/<run>.pt (12 files, generative models only): the flow-matching MLP over standardized latents (EMA weights) with the standardization statistics.
• manifest.json: size, SHA256, dataset, seed, steps, paper tables and archived metrics of every file. Total payload is about 5.9 GB.

With the companion repository:

python scripts/download_checkpoints.py --output checkpoints   # verifies SHA256
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
python evaluate.py --ckpt checkpoints/checkpoints/pg_ca_rep_l01z09.pt --out eval/pg_ca_rep_l01z09 --tasks latents,probe

evaluate.py and the data preparation scripts reproduce the paper's metrics; the datasets are not included.

Intended scope and limitations

Research on representation learning and diffusion models at low resolution (32×32 and 64×64). The models reflect their training data, including the demographic composition of CelebA and FFHQ, and are not intended for identifying people, inferring personal attributes in applications, or production image generation. The FFHQ labels used for the probes are automatically inferred attributes (age, gender, eyeglasses) from the ffhq-features-dataset and inherit its errors. Evaluation values can differ slightly across hardware, batch size and library versions.

Attribution and terms

The weights were trained on FashionMNIST, CIFAR10, CelebA, FFHQ and 3DShapes. CelebA and FFHQ are available for non-commercial research purposes only, so these weights are released under CC BY-NC 4.0. The code is MIT-licensed; see the repository's LICENSE and THIRD_PARTY_NOTICES.md.

@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}
}