MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

Welcome to the official implementation for MDNS (accepted at NeurIPS 2025), for training neural samplers for discrete distribution with masked discrete diffusion models.

Environment

conda env create -f environment.yml
conda activate mdns

Training

The commands for training MDNS on Ising and Potts models are listed in train_commands.sh

Checkpoints

We included checkpoints trained with MDNS for Ising and Potts model on 2D square lattice under the folder checkpoints.

• For Ising model, we include those on $16\times16$ and $24\times24$ square lattices, across three different inverse temperatures $\beta_{\text{high}} = 0.28$, $\beta_{\text{critical}} = 0.4407$, and $\beta_{\text{low}} = 0.6$. We named them correspondingly as ising_high.pth, ising_crit.pth and ising_low.pth under the directory.
• For Potts model, we include those on $16\times16$ square lattice with $q=3$, across three different inverse temperatures $\beta_{\text{high}} = 0.5$, $\beta_{\text{critical}} = 1.005$, and $\beta_{\text{low}} = 1.2$. We named them correspondingly as potts_high.pth, potts_crit.pth and potts_low.pth under the directory.

Evaluation

We include evaluation and visualization script of Ising and Potts model in ising_model_eval.ipynb and potts_model_eval.ipynb respectively.

Citation

If you find our work and repo help, we would appreciate your citations 🥰

@inproceedings{zhu2025mdns,
  title     = {{MDNS}: Masked Diffusion Neural Sampler via Stochastic Optimal Control},
  author    = {Zhu, Yuchen and Guo, Wei and Choi, Jaemoo and Liu, Guan-Horng and Chen, Yongxin and Tao, Molei},
  booktitle = {The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year      = {2025},
  url       = {https://openreview.net/forum?id=xIH95kXNR2}
}

Acknowledgement

Our code is partially based on the rope-vit repository (under Apache License 2.0).