NeurIPS 2025

MDNS: Masked Diffusion Neural Sampler via
Stochastic Optimal Control
Yuchen Zhu1,∗, Wei Guo1,∗, Jaemoo Choi1, Guan-Horng Liu2, Yongxin Chen1, Molei Tao1
1Georgia Institute of Technology
2FAIR at Meta
{yzhu738, wei.guo, jchoi843, yongchen, mtao}@gatech.edu, ghliu@meta.com
Abstract
We study the problem of learning a neural sampler to generate samples from
discrete state spaces where the target probability mass function π ∝e−U is
known up to a normalizing constant, which is an important task in fields such as
statistical physics, machine learning, combinatorial optimization, etc. To better
address this challenging task when the state space has a large cardinality and
the distribution is multi-modal, we propose Masked Diffusion Neural Sampler
(MDNS), a novel framework for training discrete neural samplers by aligning
two path measures through a family of learning objectives, theoretically grounded
in the stochastic optimal control of the continuous-time Markov chains. We
validate the efficiency and scalability of MDNS through extensive experiments
on various distributions with distinct statistical properties, where MDNS learns to
accurately sample from the target distributions despite the extremely high problem
dimensions and outperforms other learning-based baselines by a large margin. A
comprehensive study of ablations and extensions is also provided to demonstrate
the efficacy and potential of the proposed framework. Our code is available at
https://github.com/yuchen-zhu-zyc/MDNS.
1
Introduction
Drawing samples from an unnormalized target distribution π ∝e−U on some state space X0 is a
fundamental problem with wide-ranging applications across fields such as statistical physics [LB14],
Bayesian inference [Gel+13], machine learning [And+03; WJ08], etc. For decades, Markov chain
Monte Carlo (MCMC) methods, such as the Langevin Monte Carlo, the Metropolis-Hastings al-
gorithm, and the Glauber dynamics, have been a cornerstone. These methods simulate a Markov
chain whose ...