CoolFace
Modelpublic

chen-hao-chao/mdm-prime

sourceHugging Faceapache-2.0updated 8mo agoView on Hugging Face
1likes
Model Card

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*.

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 huggingfacehub library via `pip install -U huggingfacehub` 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.

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},
}