MDM-Prime

MDM-Prime is a discrete diffusion model enhanced with the Partial masking scheme (Prime). It enables fine-grained denoising and improves generation quality across both image and text domains. This model was proposed in our paper Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking.

---

Model Details

• Text Generation
  • Dataset: openwebtext (OWT)
  • Model Size: 92M, 286M, 375M, 860M
  • Context Length: 1,024
• Image Synthesis
  • Dataset: CIFAR-10, ImageNet-32
  • Model Size: 114M
  • Context Length: 32x32x3

---

How to Use

To download the weights, one can download the huggingface_hub library via pip install -U huggingface_hub and perform the following python code:

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="chen-hao-chao/mdm-prime",
    filename="${checkpoint_name}.pth"
)

Replace ${checkpoint_name}.pth with ${task}/${dataset}/${setup}/${checkpoint_name}.pth (e.g., image/imagenet32/results_prime_l8_imagenet32/checkpoint-599.pth). This repository is organized as follows:

mdm-prime/
├── README.md
├── image/
|   ├── cifar10/
|   └── imagenet/
|       ├── results_mdm_imagenet32/
|       ├── results_prime_supertoken_imagenet32/
|       ├── results_prime_l2_imagenet32/
|       ├── results_prime_l3_imagenet32/
|       ├── results_prime_l4_imagenet32/
|       ├── results_prime_l6_imagenet32/
|       └── results_prime_l8_imagenet32/
|            └── checkpoint-599.pth
└── text/
    └── owt/
      ├── results_prime_l2_owt/
      ├── results_prime_l2_co_owt/
      ├── results_prime_l3_owt/
      ├── results_prime_l3_co_owt/
      ├── results_prime_l4_owt/
      ├── results_prime_l4_co_owt/
      ├── results_prime_l6_owt/
      ├── results_prime_l6_co_owt/
      ├── results_prime_l8_owt/
      └── results_prime_l8_co_owt/
           └── checkpoint.ckpt

For more details regarding the training and inference processes, please refer to our github repository: chen-hao-chao/mdm-prime.

---

Citing MDM-Prime

If you find this code implementation useful, please consider citing our paper.

@inproceedings{chao2025mdmprime,
      title = {{Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking}}, 
      author = {Chen-Hao Chao, Wei-Fang Sun, Hanwen Liang, Chun-Yi Lee, Rahul G. Krishnan},
      booktitle = {Proceedings of the Conference on Neural Information Processing Systems (NeurIPS)},
      year = {2025},
}