---
license: apache-2.0
tags:
- diffusion
- discrete
- image-generation
- text-generation
library_name: transformers
inference: false
datasets:
- uoft-cs/cifar10
- jiachenlei/imagenet
- Skylion007/openwebtext
metrics:
- perplexity
base_model:
- kuleshov-group/mdlm-owt
---

# 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*](https://arxiv.org/abs/2505.18495).

---

## 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:

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

---

## Citing MDM-Prime

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

```bib
@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},
}
```