jdeschena/duo2-cifar10

Duo (Image Modeling) - CIFAR-10

This repository contains pre-trained checkpoints for image modeling on CIFAR-10, as presented in the paper The Diffusion Duality, Chapter II: $\Psi$-Samplers and Efficient Curriculum.

• Paper: The Diffusion Duality, Chapter II: $\Psi$-Samplers and Efficient Curriculum
• Project Page: s-sahoo.com/duo-ch2
• GitHub Repository: s-sahoo/duo

Model Description

Uniform-state discrete diffusion models excel at few-step generation and guidance due to their ability to self-correct. This checkpoint is part of the Duo series, which introduces a family of Predictor-Corrector (PC) samplers called $\Psi$-samplers. Unlike conventional samplers, these methods continue to improve quality as the number of sampling steps increases.

The CIFAR-10 models are trained for 1.5M steps and have approximately 35M parameters. The architecture is the same as in D3PM.

Sampling with the Duo Checkpoints

To sample from the pre-trained MDLM & Duo models, you can either play with our Colab notebook, or download the raw checkpoints from this repository, clone our GitHub repo, and run the following command:

TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1  # Depending on your PyTorch version, this might be needed to load the checkpoint
python -u -m main \
    mode=fid_eval \
    sampling.steps=64 \
    sampling.guid_weight=1.0 \
    data=cifar10 \
    data.cache_dir=<YOUR-DATA-CACHE-PATH> \
    model=unet \
    noise=cosine \
    algo=duo_base \
    algo.backbone=unet \
    loader.eval_batch_size=50 \
    eval.checkpoint_path=<PATH-TO-THE-DUO-CHECKPOINT>

Find the text checkpoints here.

Citation

If you use this work, please cite the following:

@inproceedings{
    deschenaux2026the,
    title={The Diffusion Duality, Chapter {II}: \${\textbackslash}Psi\$-Samplers and Efficient Curriculum},
    author={Justin Deschenaux and Caglar Gulcehre and Subham Sekhar Sahoo},
    booktitle={The Fourteenth International Conference on Learning Representations},
    year={2026},
    url={https://openreview.net/forum?id=RSIoYWIzaP}
}

@inproceedings{
    sahoo2025the,
    title={The Diffusion Duality},
    author={Subham Sekhar Sahoo and Justin Deschenaux and Aaron Gokaslan and Guanghan Wang and Justin T Chiu and Volodymyr Kuleshov},
    booktitle={Forty-second International Conference on Machine Learning},
    year={2025},
    url={https://openreview.net/forum?id=9P9Y8FOSOk}
}