MDM-Prime-v2-C4

MDM-Prime-v2 is an enhanced version of the MDM-Prime framework. 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 repository contains the models presented in the paper:

• MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models

Links:

• Project Page: https://chen-hao-chao.github.io/mdm-prime-v2/
• GitHub Repository: https://github.com/chen-hao-chao/mdm-prime-v2

---

Model Details

• Dataset: C4 (English)
• Model Size: 14M - 3.4B
• Context Length: 2,048

---

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-v2-c4",
    filename="${checkpoint_name}"
)

Replace ${checkpoint_name} with ${model}/${setup} (e.g., prime/prime_param_3426M_iter_14000). This repository is organized as follows:

mdm-prime-v2-c4/
├── README.md
├── arm/
├── mdm/
└── prime/
    ├── prime_param_14M_iter_70000/
    ├── prime_param_25M_iter_80000/
    ├── .../
    └── prime_param_3426M_iter_14000/
         ├── latest_checkpointed_iteration.txt
         └── iter_0014000

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

---

Citing MDM-Prime and MDM-Prime-v2

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

@article{chao2026mdmprimev2,
      title = {{MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Compute-optimal Scaling of Diffusion Language Models}}, 
      author = {Chen-Hao Chao, Wei-Fang Sun, Junwei Quan, Chun-Yi Lee, Rahul G. Krishnan},
      year = {2026},
}
@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},
}