# MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

arXiv: https://arxiv.org/abs/2508.10684 · HTML: https://mlsanity.com/paper/2508.10684/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 95 of its 128 references that are in the map.
- First posted on arXiv 14 Aug 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 95)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Pyramidal Flow Matching for Efficient Video Generative Modeling](https://mlsanity.com/paper/2410.05954/) (2024)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Scaling Diffusion Language Models via Adaptation from Autoregressive Models](https://mlsanity.com/paper/2410.17891/) (2024)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions](https://mlsanity.com/paper/2502.06768/) (2025)
- [Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling](https://mlsanity.com/paper/2409.02908/) (2024)
- [Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations](https://mlsanity.com/paper/1706.04702/) (2017)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)

## Built on by (most cited first, 21 of 21)

- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models](https://mlsanity.com/paper/2509.21655/) (2025)
- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [Proximal Diffusion Neural Sampler](https://mlsanity.com/paper/2510.03824/) (2025)
- [TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://mlsanity.com/paper/2509.25171/) (2025)
- [Discrete Adjoint Matching](https://mlsanity.com/paper/2602.07132/) (2026)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddings](https://mlsanity.com/paper/2509.22925/) (2025)
- [Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion](https://mlsanity.com/paper/2605.23346/) (2026)
- [Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces](https://mlsanity.com/paper/2511.06239/) (2025)
- [Discrete Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2602.08243/) (2026)
- [Rare Event Analysis via Stochastic Optimal Control](https://mlsanity.com/paper/2604.13213/) (2026)
- [TABES: Trajectory-Aware Backward-on-Entropy Steering for Masked Diffusion Models](https://mlsanity.com/paper/2602.00250/) (2026)
- [Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spaces](https://mlsanity.com/paper/2602.05961/) (2026)
- [Entangled Schr"odinger Bridge Matching](https://mlsanity.com/paper/2511.07406/) (2025)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [Efficient Inference for Coupled Hidden Markov Models in Continuous Time and Discrete Space](https://mlsanity.com/paper/2510.12916/) (2025)
- [A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching](https://mlsanity.com/paper/2606.24140/) (2026)
- [Discrete Diffusion Models via Evolving Variational Autoregressive Networks](https://mlsanity.com/paper/2609.27306/) (2026)
- [Iterative Exact Discrete Guidance for Energy-Based Sampling](https://mlsanity.com/paper/2609.37043/) (2026)
- [Fixed-point neural samplers on discrete spaces](https://mlsanity.com/paper/2610.01739/) (2026)

## Abstract (arXiv)

We study the problem of learning a neural sampler to generate samples from discrete state spaces where the target probability mass function πproptoe-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.