---
datasets:
- allenai/c4
license: apache-2.0
metrics:
- perplexity
pipeline_tag: text-generation
---

# 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 Scaling of Diffusion Language Models](https://huggingface.co/papers/2603.16077)

**Links:**
- **Project Page:** [https://chen-hao-chao.github.io/mdm-prime-v2/](https://chen-hao-chao.github.io/mdm-prime-v2/)
- **GitHub Repository:** [https://github.com/chen-hao-chao/mdm-prime-v2](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:

```python
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](https://github.com/chen-hao-chao/mdm-prime-v2).

---

## Citing MDM-Prime and MDM-Prime-v2

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

```bibtex
@article{chao2026mdmprimev2,
      title = {{MDM-Prime-v2: Binary Encoding and Index Shuffling Enable 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},
}
```