Title: MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
Venue: NeurIPS 2025
Accepted: yes
Keywords: Masked diffusion model

Continuous-time Markov Chain

Sampling

Markov-chain Monte Carlo

Optimal Control
Abstract: We study the problem of learning a neural sampler to generate samples from discrete state spaces where the target probability mass function $\pi\propto\mathrm{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 **M**asked **D**iffusion **N**eural **S**ampler (**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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to leverage a Stochastic Optimal Control (SOC) of Continuous Time Markov Chains (CTMC) for solving the problem of sampling from an unnormalized density  $\\pi(x) \\propto exp(U(x))$ on a discrete state space. The authors state that sampling from an unnormalized density problem can be formulated as an SOC problem and propose computationally efficient methods to solve the corresponding SOC instance of the problem. In contrast to the previous approach for solving the SOC of CTMC, the authors eliminate differentiation through the sample generation trajectory and provide a total of 4 optimization objectives, i.e., REINFORCE trick-based, Log Variance Divergence-based, two Cross Entropy-based, that require only trajectory samples, but do not differentiate through them. \n\nThe algorithm for solving the SOC of CTMC for sampling with a Masked (and Uniform in Appendix F) prior process is proposed, called MNDS (and UNDS), and tested on a series of statistical physics-based benchmarks. The method is compared to the state-of-the-art neural sampler LEAPS [1] and learning-free Metropolis Hastings baseline, and shows itself on par on high temperature distributions and superior on lower temperature distributions.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- The proposed approach is quite simple and effective.\n- SOC for CTMC development helps to bridge the gap between CTMC models and diffusion models and, in my opinion, is valuable for the CTMC community.\n- The MNDS method seems to outperform other methods on lower temperature distributions.\n- Theoretical guarantees on sampling quality and normalizing constant are shown.\n- The paper is mainly well-written and easy to follow.\n- Experimental evaluation is through with a lot of quantitative metrics, which is rare for an unnormalized density on a discrete state space sampler paper.\n\nWeaknesses:\n\n- In Section 3.2, when LogVar and CE losses are presented, it is hard to understand where exactly the stop gradient can be placed and why it is so. I think it should be explained more clearly.\n- Would like to see the code in the sake of reproducibility, although sufficient experimental details are described in the Appendix.\n- No information on the computational complexity of the method compared to LEAPS and MH. It would be nice if the authors provided computational complexity analysis and training/inference time comparisons.\n- Performance seems a bit behind the MH on Ising 4x4, see Table 2.\n- I can’t find the preconditioning strategy mentioned in the main text, only in the Appendix. I think it should be mentioned in the Sec 4.\n- The model was evaluated only on the statistical physics benchmarks. It would be nice to see it’s performance on other type of problems, e.g., combinatorial optimization.\n\nQUESTIONS:\n- It would be nice if the authors added LEAPS to the comparison in theTable 2.\n- Why is WDCE loss used for experiments on 16x16 Ising and Potts models, while the LV loss showed itself to be the best on the training objectives study, see Sec 4.1?\n- In the Proposition 2 main paper, there seems to be a missing expectation. I think it should be $\\hat{Z} = \\mathbb{E}_{X \\sim P^u}e^{W^u(x)}$\n- How do you ensure the memoryless condition $\\int_0^T \\gamma(t)dt=\\infty$ in practice? In the case of masked process MNDS memoryless property seems obvious (since you start from a point), but for the UDLM (Appendix F), in my opinion, it is not obvious how to ensure this condition.\n\nLIMITATIONS:\nWould like to see the mention of method not being simulation-free in the limitations section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the task of training a neural sampler without data, on discrete distributions by only using the energy function. This paper combines many approaches in continuous space and applies them to CTMC. This paper also proposed an efficient training objection that does not require simulation every gradient update. The authors provided comprehensive evaluation to demonstrate the efficientness of the proposed approach.\n\nSTRENGTHS AND WEAKNESSES:\n## Strength:\n\n1. The writing is clear, pleasant to read, and easy to follow. It is a pleasure to review and read this paper.\n\n2. The paper not only proposed one approach. Instead, it provides a systematic recipe with several loss choices one can consider. While most of  these are not novel, it makes the paper complete.\n\n3. The proposed simulation-free loss is quite interesting and effective. \n\n4. Experiments are convincing. Comparing different losses is also helpful. \n\n\n\n## Weakness:\n\n1. I might suggest adding citations and discussion on previous works that use similar ideas to the buffer and weighted loss. As far as I know, FAB[1] and Adjoint Sampler[2]'s buffers have similar ideas. In FAB, one can do IS to reweight samples from the buffer. In AS, one collects buffer using current control and train the network with TSM.\n\n2. By convention, one define potential /  Energy $U$ that $p \\propto \\exp(-U)$. However, in this paper, it was defined as $p \\propto \\exp(U)$ andm hence it might be beneficial to follow that notation. I do not think this is a weakness (but I do not have other places to put this).\n\n\n\n[1] Midgley, Laurence Illing, et al. \"Flow Annealed Importance Sampling Bootstrap.\" ICLR,\n\n[2] Havens, Aaron, et al. \"Adjoint sampling: Highly scalable diffusion samplers via adjoint matching.\" ICML.\n\nQUESTIONS:\n1. I am not generally familiar with the setup of these systems. Therefore, I might be curious why some of the parameters are set in that way (for example, line 301)? Is it to make the system easy (like reducing phase transitions) or make it harder? Is it a convention?\n\n\n2. I am a bit lost in the equation below line 202 (is there also a typo?). Could you please explain this to me?\n\nLIMITATIONS:\nYes,  the authors have adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-25"}

comments: SUMMARY:
This paper extends stochastic optimal control for continuous-time Markov chains (CTMCs) to discrete state spaces, introducing the Masked Diffusion Neural Sampler (MDNS)-–a neural sampler tailored to discrete domains. It proposes a weighted denoising cross-entropy loss, a scalable objective for training masked-diffusion models in sampling tasks. The theoretical framework generalizes existing continuous-space neural samplers to discrete settings, while a replay-buffer–based training algorithm improves computational efficiency. Experiments are comprehensive, benchmarking MDNS against both learning-based and classical non-learning Markov-chain methods across diverse probabilistic metrics.

STRENGTHS AND WEAKNESSES:
**Strengths**

1. Strong theory and a solid contribution that extends CTMC-based stochastic optimal control to discrete spaces.
2. Introduces an efficient loss function and accompanying algorithm.

**Weaknesses**

1. The paper does not report how trajectory-level loss computation affects the variance of learning.
2. The importance weights may further amplify training variance, yet no variance-reduction estimator is provided.
3. Motivation is thin: why adopt a CTMC framework in discrete spaces when established approaches like GFlowNets already serve as effective discrete samplers without relying on CTMC dynamics?

QUESTIONS:
1. **Replay buffer**: Does it materially improve training? Have you explored prioritized replay?
2. **Trajectory length**: Can the algorithm scale to very long trajectories?
3. **CTMC justification**: Why is a CTMC formulation needed in discrete spaces?
4. **Variance control**: What techniques could you use to reduce training variance?

LIMITATIONS:
No
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 4
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new diffusion sampler for discrete problems that is based on optimal control of a continuous-time Markov chain (CTMC). To do so, the authors first show that we can obtain a generator by minimizing the KL divergence between the optimal path distribution from a masked diffusion sampler and the generator. The authors discuss and compare different loss functions including REINFORCE loss, log-variance loss and cross entropy loss to solve the SOC problem. finally, they introduce a new weigthed cross entropy loss that is supposely computationally more efficient as it can reuse data. The authors compare their algorithm on challenging Ising models and potts models to a recent SOTA method (LEAPS) and show that it outperforms this SOTA method.\n\nSTRENGTHS AND WEAKNESSES:\nstrength:\n- interesting new algorithm for discrete sampling\n- sound derivation (even though I could not fully follow everything) based on stochastic optimal control\n- good experimental results with a clear improvement over SOTA on challenging problems:\n\nweaknesses:\n- clarity and motivation. The paper is hard to follow, notation is hard to grasp and it is unclear why certain things are introduced (e.g. Lemma 1). Moreover, basic definitions are hidden in equations and are quite hard to find (e.g. the defintion of W^u, that should be an own equation with W^u = ...) and no intution is given what this terms really mean. \n- Results are not fully conclusive (see questions).\n\nQUESTIONS:\n- It should be clarified how the optimal control formulation for discrete samplers deviates from the standard optimal control formulation for continuous samplers\n- From the study on the simpler 4x4 task, it seems that the log-variance loss is performing best. However, the paper proceeds with the weighted cross entropy loss (WCE). The advantage of the WCE to logVariance is unclear. Is it computation time due to data reuse? If so, that need to be evaluated as well\n- Its unclear to me why we can not implement a data-reuse mechanism (e.g. with a similar importance weighting) for the other loss functions. Log-Variance loss seems to be better then cross entropy, so  the motivation why this was only done for the cross entropy loss is missing for me\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}