MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

2025-08-14

AI Summary

This work addresses efficient sampling from target distributions π ∝ e⁻ᵁ with unknown normalizing constants over discrete state spaces—common in statistical physics and combinatorial optimization—where distributions are high-dimensional and strongly multimodal. We propose the Masked Diffusion Neural Sampler (MDNS), grounded in stochastic optimal control theory, which frames sampling as path-measure alignment on continuous-time Markov chains. MDNS introduces a learnable masked diffusion mechanism and a family of path-alignment objective functions. Compared to conventional MCMC and existing learning-based samplers, MDNS achieves significantly improved sampling accuracy and faster convergence on high-dimensional, large-scale, and highly multimodal distributions, while exhibiting strong scalability and cross-domain applicability. Extensive experiments demonstrate that MDNS outperforms state-of-the-art baselines across multiple benchmark tasks, including spin-glass models, Ising inference, and combinatorial optimization instances.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We study the problem of learning a neural sampler to generate samples from discrete state spaces where the target probability mass function $πproptomathrm{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 $ extbf{M}$asked $ extbf{D}$iffusion $ extbf{N}$eural $ extbf{S}$ampler ($ extbf{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.

Problem

• Learning neural sampler for discrete state spaces
• Handling large cardinality and multi-modal distributions
• Efficient sampling from high-dimensional target distributions

Innovation

• Masked Diffusion Neural Sampler framework
• Stochastic optimal control for Markov chains
• Learning objectives aligning path measures