Proximal Diffusion Neural Sampler

Technical Summary: Proximal Diffusion Neural Sampler (PDNS)

Problem Statement
Learning a diffusion-based neural sampler to draw samples from an unnormalized target distribution (e.g., a Boltzmann distribution $\pi(x) \propto e^{-\beta E(x)}$) is a fundamental task in computational statistics, Bayesian inference, and statistical physics. This problem is often formulated as a Stochastic Optimal Control (SOC) problem on path measures, where a neural network parameterizes a control process to transport a simple reference distribution to the target.

However, training these samplers faces significant challenges when the target distribution is multimodal with high energy barriers separating the modes. In such scenarios, standard SOC-based methods (including Cross-Entropy and Weighted Denoising Cross-Entropy variants) often suffer from mode collapse. Early in training, the model captures only a few modes; the global loss function then reinforces these captured modes based on the limited trajectories generated, causing the probability mass to concentrate on a subset of modes while ignoring others. This leads to unstable optimization and poor coverage of the target distribution.

Methodology
The authors propose the Proximal Diffusion Neural Sampler (PDNS), a unified framework designed to address mode collapse by reformulating the training process as a sequence of local sub-problems rather than a single global minimization.

1. Unified Framework: The paper first unifies continuous and discrete diffusion samplers under a single SOC framework. It defines the optimal path measure $P^*$ as the solution to minimizing a relative entropy objective subject to a reference path measure $P_{ref}$ that satisfies a "memoryless" condition (independence between initial and terminal distributions).
2. Proximal Point Method: To mitigate mode collapse, PDNS applies the proximal point method to the space of path measures. Instead of directly optimizing the global objective, the algorithm iteratively solves a local sub-problem at step $k$:
   $$P_{\theta^*_k} = \arg\min_{P_\theta} \left( -E_{P_\theta}[r(X_T)] + KL(P_\theta \| P_{ref}) + \frac{1}{\eta_k} KL(P_\theta \| P_{\theta_{k-1}}) \right)$$
   Here, $\eta_k$ is a proximal step size, and the term $KL(P_\theta \| P_{\theta_{k-1}})$ acts as a regularizer keeping the current iterate close to the previous one.
3. Theoretical Properties: The authors prove that the optimal solution to the $k$-th proximal step is a geometric interpolation between the previous iterate and the optimal path measure. Specifically, the terminal distribution of the $k$-th step, $P_k$, is a geometric interpolation between the reference terminal distribution $\nu$ and the target $\pi$. As $k \to \infty$, the sequence of path measures converges to the target $P^*$. This staged procedure ensures that the model explores the state space gradually, preventing premature commitment to a single mode.
4. Practical Instantiation (Proximal WDCE): For efficiency, the authors instantiate PDNS using a Proximal Weighted Denoising Cross-Entropy (Proximal WDCE) objective. This approach avoids storing full trajectories by using importance weights accumulated online.
   • Continuous Case: The objective involves matching a denoising score (or bridge) with a weight term derived from the proximal step size $\eta_k$ and the energy function.
   • Discrete Case: The objective adapts masked discrete diffusion, utilizing a similar weighting scheme based on the Radon-Nikodym derivative of the path measures.
5. Scheduler Design: The paper discusses strategies for selecting the step size $\eta_k$ (or its equivalent $\lambda_k$). A smaller $\eta_k$ (stronger regularization) preserves mode coverage but slows convergence, while a larger $\eta_k$ risks mode collapse. The authors propose both predefined schedules (linearly decreasing $\lambda_k$) and an adaptive strategy that adjusts $\eta_k$ based on a trust-region constraint on the KL divergence between the current model and the next target.

Key Contributions

1. Unified Framework: The paper integrates continuous and discrete SOC-based samplers into a single path measure formulation.
2. Proximal Diffusion Neural Sampler (PDNS): It introduces a novel framework that leverages proximal point iterations in path space to stabilize learning and mitigate mode collapse in multimodal settings.
3. Proximal WDCE: It derives concrete, computationally efficient instantiations of PDNS for both continuous and discrete domains using proximal weighted denoising cross-entropy losses.
4. Scheduler Strategies: The authors outline principled strategies for selecting proximal step sizes, including an adaptive scheduler that balances convergence speed and stability.

Experimental Results
The authors evaluate PDNS on a wide range of continuous and discrete benchmarks:

• Continuous Tasks:
  • Synthetic Energy Functions: PDNS outperforms strong baselines (including SMC, DDS, PIS, and AS) on multimodal benchmarks like the 32-well potential (MW54), Funnel, and high-dimensional Gaussian mixtures (GMM40). It achieves state-of-the-art results on the Student's t-mixture (MoS) and particle-based potentials (Lennard-Jones systems LJ-13 and LJ-55).
  • Molecular Dynamics: On the Alanine Dipeptide (AD) system, PDNS achieves KL divergences comparable to the state-of-the-art ASBS method and accurately captures conformational structures in Ramachandran plots.
• Discrete Tasks:
  • Statistical Physics: PDNS significantly outperforms LEAPS and Metropolis-Hastings (MH) in sampling Ising and Potts models at low temperatures (near or below critical points), where mode separation is severe. It achieves lower errors in magnetization and 2-point correlations.
  • Combinatorial Optimization: Applied to the Maximum Cut problem on Barabási–Albert and Erdős–Rényi graphs, PDNS demonstrates the ability to reach solution qualities comparable to ground-truth solvers (Gurobi).

Significance and Claims
The paper claims that PDNS provides a robust solution to the mode collapse problem inherent in training diffusion-based neural samplers for complex, multimodal distributions. By decomposing the global optimization into a sequence of proximal steps, the method promotes thorough exploration across modes and stabilizes optimization in high-dimensional spaces.

The authors note two main limitations:

1. The investigation focuses on a single instantiation of PDNS (proximal WDCE), leaving room for exploring other proximal objectives or priors.
2. Experiments are conducted on controlled benchmarks rather than large-scale real-world applications.

The work suggests future directions include exploring SOC problems with non-memoryless reference dynamics, extending the framework to more general state spaces, and developing other instantiations of the proximal objective.