mlsanityData to 14 Oct 2026

MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

arXiv 2508.10684 ↗First posted 14 Aug 2025cs.LG math.OC stat.COcited by 21

On the map: Diffusion models → Diffusion foundations and sampling

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.

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.

Builds on

95 of 128 references.

Closest ancestors · 25

Same direction, newest first

  1. 2025Improving Classifier-Free Guidance in Masked Diffusion: Low-Dim Theoretical Insights with High-Dim Impact11 Jul
  2. 2025Adjoint Schr"odinger Bridge Sampler27 Jun
  3. 2025Non-equilibrium Annealed Adjoint Sampler22 Jun
  4. 2025What Exactly Does Guidance Do in Masked Discrete Diffusion Models12 Jun
  5. 2025Progressive Tempering Sampler with Diffusion5 Jun
  6. 2025Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics12 Feb
  7. 2025Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms1 Feb
  8. 2024Simple Guidance Mechanisms for Discrete Diffusion Models13 Dec
  9. 2024Diffusion-PINN Sampler20 Oct
  10. 2024textitJump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models10 Oct
  11. 2024How Discrete and Continuous Diffusion Meet: Comprehensive Analysis of Discrete Diffusion Models via a Stochastic Integral Framework4 Oct
  12. 2024A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization3 Jun
13 more
  1. 2024Scaling Rectified Flow Transformers for High-Resolution Image Synthesis5 Mar
  2. 2024Zeroth-Order Sampling Methods for Non-Log-Concave Distributions: Alleviating Metastability by Denoising Diffusion27 Feb
  3. 2024Convergence Analysis of Discrete Diffusion Model: Exact Implementation through Uniformization12 Feb
  4. 2023Denoising Diffusion Samplers27 Feb
  5. 2022Score-based Continuous-time Discrete Diffusion Models30 Nov
  6. 2022From Denoising Diffusions to Denoising Markov Models7 Nov
  7. 2022An optimal control perspective on diffusion-based generative modeling2 Nov
  8. 2022Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning8 Aug
  9. 2022A Continuous Time Framework for Discrete Denoising Models30 May
  10. 2021Structured Denoising Diffusion Models in Discrete State-Spaces7 Jul
  11. 2020Score-Based Generative Modeling through Stochastic Differential Equations26 Nov
  12. 2020Denoising Diffusion Implicit Models6 Oct
  13. 2020Denoising Diffusion Probabilistic Models19 Jun

Other references · 70

  1. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  2. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,058
  3. 2020Training data-efficient image transformers & distillation through attentioncited by 4,586
  4. 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusioncited by 2,838
  5. 2016Categorical Reparameterization with Gumbel-Softmaxcited by 2,739
  6. 2022MaskGIT: Masked Generative Image Transformercited by 1,099
  7. 2025Large Language Diffusion Modelscited by 912
  8. 2024Simple and Effective Masked Diffusion Language Modelscited by 848
  9. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distributioncited by 745
  10. 2024Simplified and Generalized Masked Diffusion for Discrete Datacited by 626
  11. 2017Solving high-dimensional partial differential equations using deep learningcited by 556
  12. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 409
38 more
  1. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Datacited by 381
  2. 2024Pyramidal Flow Matching for Efficient Video Generative Modelingcited by 366
  3. 2013Black Box Variational Inferencecited by 319
  4. 2024Discrete Flow Matchingcited by 317
  5. 2024Scaling Diffusion Language Models via Adaptation from Autoregressive Modelscited by 274
  6. 2024Rotary Position Embedding for Vision Transformercited by 236
  7. 2024Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Controlcited by 233
  8. 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusionscited by 214
  9. 2024Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Samplingcited by 208
  10. 2017Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equationscited by 191
  11. 2025Remasking Discrete Diffusion Models with Inference-Time Scalingcited by 189
  12. 2021Path Integral Sampler: a stochastic control approach for samplingcited by 181
  13. 2023Protein Design with Guided Discrete Diffusioncited by 157
  14. 2023A Reparameterized Discrete Diffusion Model for Text Generationcited by 152
  15. 2024Unlocking Guidance for Discrete State-Space Diffusion and Flow Modelscited by 129
  16. 2022Flow Annealed Importance Sampling Bootstrapcited by 127
  17. 2014Neural Variational Inference and Learning in Belief Networkscited by 117
  18. 2024Iterated Denoising Energy Matching for Sampling from Boltzmann Densitiescited by 116
  19. 2024Diffusion Language Models Are Versatile Protein Learnerscited by 98
  20. 2025Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matchingcited by 85
  21. 2020Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path spacecited by 82
  22. 2024Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Designcited by 79
  23. 2024Beyond Autoregression: Fast LLMs via Self-Distillation Through Timecited by 78
  24. 2024Discrete-state Continuous-time Diffusion for Graph Generationcited by 76
  25. 2024NETS: A Non-Equilibrium Transport Samplercited by 69
  26. 2024Think While You Generate: Discrete Diffusion with Planned Denoisingcited by 63
  27. 2024Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspectivecited by 57
  28. 2021Bayesian Learning via Neural Schr"odinger-F"ollmer Flowscited by 56
  29. 2024Meissonic: Revitalizing Masked Generative Transformers for Efficient High-Resolution Text-to-Image Synthesiscited by 56
  30. 2025Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Modelcited by 55
  31. 2022Score-Based Diffusion meets Annealed Importance Samplingcited by 51
  32. 2024Sequential Controlled Langevin Diffusionscited by 48
  33. 2025Theoretical Benefit and Limitation of Diffusion Language Modelcited by 47
  34. 2023Learning Interpolations between Boltzmann Densitiescited by 43
  35. 2024PepTune: De Novo Generation of Therapeutic Peptides with Multi-Objective-Guided Discrete Diffusioncited by 37
  36. 2025Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Imagescited by 37
  37. 2023Bridging Discrete and Backpropagation: Straight-Through and Beyondcited by 31
  38. 2025LEAPS: A discrete neural sampler via locally equivariant networkscited by 31

50 most cited of 70 All 70 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 9; 2026: 12.

21 papers.

Sort:

By area:

  1. 2025Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inferencecited by 23
  2. 2025DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Modelscited by 23
  3. 2025Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimizationcited by 20
  4. 2025Proximal Diffusion Neural Samplercited by 16
  5. 2025TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusioncited by 14
  6. 2026Discrete Adjoint Matchingcited by 11
  7. 2026Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matchingcited by 10
  8. 2025Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddingscited by 6
  9. 2026Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusioncited by 5
  10. 2025Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spacescited by 4
  11. 2026Discrete Adjoint Schr"odinger Bridge Samplercited by 4
  12. 2026Rare Event Analysis via Stochastic Optimal Controlcited by 4
9 more
  1. 2026TABES: Trajectory-Aware Backward-on-Entropy Steering for Masked Diffusion Modelscited by 3
  2. 2026Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spacescited by 2
  3. 2025Entangled Schr"odinger Bridge Matchingcited by 1
  4. 2026FluxLite: Inference-Time Proposal Control for Discrete Diffusion Modelscited by 1
  5. 2025Efficient Inference for Coupled Hidden Markov Models in Continuous Time and Discrete Spacecited by 0
  6. 2026A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matchingcited by 0
  7. 2026Discrete Diffusion Models via Evolving Variational Autoregressive Networkscited by 0
  8. 2026Iterative Exact Discrete Guidance for Energy-Based Samplingcited by 0
  9. 2026Fixed-point neural samplers on discrete spacescited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Masked Diffusion Models are Secretly Learned-Order Autoregressive Models
  2. 2025Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
  3. 2026Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
  4. 2026From Atoms to Entropy: Optimal Noise Allocation for Diffusion Training in the Convex Regime
  5. 2026NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization
  6. 2025Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations
  7. 2025From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training
  8. 2024Discrete generative diffusion models without stochastic differential equations: a tensor network approach