🏃‍➡️ Coming Soon

Thank you for your interest in our project! We are currently preparing and testing the code and will be releasing it soon.

What will be Included

This repository contains the code implementation of the experiments presented in the paper Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking.

• 🐳 Docker support for enhanced reproducibility
• 💿 Pretrained weights for inference and evaluation
• 🔍 Sampling process visualizations to understand model behavior
• 🔬 Code for all experiments in our paper:
  • Toy experiments on synthetic data
  • Text generation on OpenWebText
  • Image generation on CIFAR-10 & ImageNet-32

Overview

Toy Examples

• Dataset: 2D Synthetic Dataset
• Folder: mdm-prime/toy
• Demo (click me)

  

Text Generation

• Dataset: OpenWebText (OWT)
• Folder: mdm-prime/text
• Demo (click me)

  

Image Generation

• Dataset: CIFAR-10, ImageNet-32
• Folder: mdm-prime/image
• Demo (click me)

  

License

This code implementation is developed based on the following repositories.

• kuleshov-group/mdlm (at commit 3ecb6dc), licensed under the Apache-2.0 license.
• facebookresearch/flow_matching (at commit 47c4396), licensed under the CC BY-NC 4.0 license.

Further changes based on this repository are licensed under the Apache-2.0 and CC BY-NC 4.0 licenses.

Citing MDM-Prime

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

@article{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},
      journal={\tt arXiv:2505.18495 [cs.LG]},
      year={2025},
}