# Proximal Diffusion Neural Sampler

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 81 references that are in the map.
- First posted on arXiv 4 Oct 2025.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (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)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data](https://mlsanity.com/paper/2406.03736/) (2024)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [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)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics](https://mlsanity.com/paper/2302.01170/) (2023)
- [Particle Denoising Diffusion Sampler](https://mlsanity.com/paper/2402.06320/) (2024)
- [A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization](https://mlsanity.com/paper/2406.01661/) (2024)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Target Score Matching](https://mlsanity.com/paper/2402.08667/) (2024)
- [How Discrete and Continuous Diffusion Meet: Comprehensive Analysis of Discrete Diffusion Models via a Stochastic Integral Framework](https://mlsanity.com/paper/2410.03601/) (2024)

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

- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [Learning Boltzmann Generators via Constrained Mass Transport](https://mlsanity.com/paper/2510.18460/) (2025)
- [Scaling Autoregressive Models for Lattice Thermodynamics](https://mlsanity.com/paper/2603.14695/) (2026)
- [Discrete Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2602.08243/) (2026)
- [Efficient Adjoint Matching for Fine-tuning Diffusion Models](https://mlsanity.com/paper/2605.11480/) (2026)
- [QUATRO: Query-Adaptive Trust Region Policy Optimization for LLM Fine-tuning](https://mlsanity.com/paper/2602.04620/) (2026)
- [Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spaces](https://mlsanity.com/paper/2602.05961/) (2026)
- [Exact Global MCMC with Denoising Diffusion](https://mlsanity.com/paper/2609.00279/) (2026)
- [A Stabilized Path-Space Approach to Diffusion-Based Posterior Sampling](https://mlsanity.com/paper/2606.12710/) (2026)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [Rethinking the Flow-Based Gradual Domain Adaptation: A Semi-Dual Optimal Transport Perspective](https://mlsanity.com/paper/2602.01179/) (2026)
- [Slack More, Predict Better: Proximal Relaxation for Probabilistic Latent Variable Model-based Soft Sensors](https://mlsanity.com/paper/2603.11473/) (2026)
- [Mode Collapse Is Cheap to Detect: A Ground-Truth-Free Pre-Flight Check for Neural Samplers](https://mlsanity.com/paper/2609.26272/) (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)

The task of learning a diffusion-based neural sampler for drawing samples from an unnormalized target distribution can be viewed as a stochastic optimal control problem on path measures. However, the training of neural samplers can be challenging when the target distribution is multimodal with significant barriers separating the modes, potentially leading to mode collapse. We propose a framework named Proximal Diffusion Neural Sampler (PDNS) that addresses these challenges by tackling the stochastic optimal control problem via proximal point method on the space of path measures. PDNS decomposes the learning process into a series of simpler subproblems that create a path gradually approaching the desired distribution. This staged procedure traces a progressively refined path to the desired distribution and promotes thorough exploration across modes. For a practical and efficient realization, we instantiate each proximal step with a proximal weighted denoising cross-entropy (WDCE) objective. We demonstrate the effectiveness and robustness of PDNS through extensive experiments on both continuous and discrete sampling tasks, including challenging scenarios in molecular dynamics and statistical physics. Our code is available at https://github.com/AlexandreGUO2001/PDNS.