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 first proposed in our paper Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking.

---

Model Details

• Text Generation
  • Dataset: openwebtext (OWT)
  • Model Size: 131M
  • 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.

@article{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},
      journal={\tt arXiv:2505.18495 [cs.LG]},
      year={2025},
}