Research map / Image generation / Fast diffusion sampling
MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MDNS, a discrete neural sampler trained via stochastic optimal control of CTMCs with masked diffusion as reference, using gradient-free objectives including a weighted denoising cross-entropy (WDCE) loss, evaluated on Ising and Potts models.
Problem
Learning neural samplers that draw samples from unnormalized discrete distributions (e.g., Ising/Potts models) with large state spaces and multimodality, where CTMC trajectories make direct gradient-based training non-differentiable.
Approach
Cast discrete sampling as SOC over CTMCs with a memoryless masked-diffusion reference path measure, then train the controlled generator with gradient-free path-measure objectives (RERF, LV, CE) and a scalable weighted denoising cross-entropy loss that importance-weights final samples and applies masked denoising supervision.
Key results
- On 16x16 Ising at beta_critical, MDNS achieves magnetization error 3.7e-3 and ESS 0.933 vs LEAPS 7.4e-3 and 0.384.
- On 16x16 Potts at beta_low=1.2, MDNS reaches correlation error 8.8e-5 and ESS 0.933 while LEAPS gets 2.5e-1 correlation error and ESS 0.012.
- On 4x4 Ising, the LV objective ranks best (ESS 0.9713, TV 0.0748), with the ranking LV > WDCE > RERF > CE on most metrics.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use MDNS instead of LEAPS | 16×16 Ising · β=0.6 | Corr. error: 5.8\mathrm{e}{-1} -> 2.4\mathrm{e}{-3} | Helped |
| use MDNS instead of LEAPS | 16×16 Ising · β=0.4407 | Corr. error: 1.6\mathrm{e}{-1} -> 2.0e{-3} | Helped |
| use MDNS instead of LEAPS | 16×16 Ising · β=0.28 | Mag. error: 7.4\mathrm{e}{-3} -> 8.5\mathrm{e}{-3} | Hurt |
| use MDNS instead of LEAPS | 16×16 Potts · β=1.2 | Mag. error: 2.9\mathrm{e}{-1} -> 1.3e{-3} | Helped |
| use MDNS instead of LEAPS | 16×16 Potts · β=1.005 | Mag. error: 2.7\mathrm{e}{-1} -> 4.3e{-3} | Helped |
| use MDNS instead of LEAPS | 16×16 Potts · β=0.5 | Mag. error: 2.9\mathrm{e}{-3} -> 2.2\mathrm{e}{-3} | Helped |
| use LV rather than CE objective (ablation) | MDNS · 4×4 Ising, β=0.28 | KL(p̂‖π): 0.0393 -> 0.0348 | Helped |
Surprising findings
- Metropolis-Hastings cannot mix on 16x16 Potts even after more than 20 hours of simulation, while it works on all Ising experiments.
- The Gumbel-softmax relaxation approach fails to converge to the correct target even in low dimensions.
Limitations
- Evaluated only on Ising and Potts lattice models; performance on graph or combinatorial optimization distributions is unknown.
- Warm-up annealing schedule is heuristic and not systematically studied.
- LV/RERF are not valid loss functions; reduction in objective does not guarantee improved performance.
Experimental setup
Datasets: Ising model samples (self-generated via target energy) Potts model samples (self-generated via target energy)
Benchmarks: 2D Ising model 4x4 and 16x16 2D Potts model 16x16 (q=3)
Metrics: ESS total variation KL divergence chi-squared divergence magnetization error 2-point correlation error absolute error of log Z
Baselines: LEAPS Metropolis-Hastings Gumbel-softmax relaxation objective
Builds on: LEAPS Path Integral Sampler denoising cross-entropy masked diffusion random order autoregressive sampler
Closest papers by meaning
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.90)
- Discrete Neural Flow Samplers with Locally Equivariant Transformer (NeurIPS 2025 · similarity 0.89)
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.88)
- LEAPS: A discrete neural sampler via locally equivariant networks (ICML 2025 · similarity 0.87)
- Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.87)
- Learning to Sample From Diffusion Models via Inverse Reinforcement Learning (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
