Proximal Diffusion Neural Sampler

2025-10-04

AI Summary

To address mode collapse in sampling from multimodal unnormalized target distributions—caused by high energy barriers—this paper proposes a diffusion-based neural sampling framework grounded in proximal point methods on the space of probability paths. The method decomposes complex sampling learning into a sequence of progressive proximal subproblems, substantially enhancing cross-modal exploration. Furthermore, we introduce a proximal weighted denoising cross-entropy (WDCE) objective, integrating diffusion modeling with stochastic optimal control theory to ensure robust training. Experiments demonstrate that our approach outperforms existing baselines on both continuous and discrete sampling tasks. Notably, in strongly multimodal domains such as molecular dynamics and statistical physics, it achieves significant improvements in distribution coverage and sampling stability.

Technology Category

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

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems
• Web Mining and Content Analysis: Content-based information diffusion

Abstract

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 extbf{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.

Problem

• Learning diffusion neural samplers for multimodal distributions with mode collapse
• Solving stochastic optimal control via proximal methods on path measures
• Enabling thorough mode exploration in continuous and discrete sampling tasks

Innovation

• Proximal point method optimizes path measures
• Staged subproblems gradually approach target distribution
• Proximal weighted denoising cross-entropy enables efficient realization