File size: 2,999 Bytes
ed9b5c4 bbc078f ed9b5c4 2b7a5da ed9b5c4 142e39d ed9b5c4 142e39d 5c7687d ed9b5c4 142e39d ed9b5c4 142e39d ed9b5c4 9facb1b ed9b5c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | ---
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},
}
```
|