Alternating Diffusion for Proximal Sampling with Zeroth Order Queries
Authors: Hirohane Takagi, Atsushi Nitanda
Organizations: Graduate School of Information Science and Technology, The University of Tokyo, Japan · Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore · Centre for Frontier AI Research, Agency for Science, Technology and Research, Singapore · College of Computing and Data Science, Nanyang Technological University, Singapore
This work introduces a new approximate proximal sampler that operates solely with zeroth-order information of the potential function. Prior theoretical analyses have revealed that proximal sampling corresponds to alternating forward and backward iterations of the heat flow. The backward step was originally implemented by rejection sampling, whereas we directly simulate the dynamics. Unlike diffusion-based sampling methods that estimate scores via learned models or by invoking auxiliary samplers, our method treats the intermediate particle distribution as a Gaussian mixture, thereby yielding a Monte Carlo score estimator from directly samplable distributions. Theoretically, when the score estimation error is sufficiently controlled, our method inherits the exponential convergence of proximal sampling under isoperimetric conditions on the target distribution. In practice, the algorithm avoids rejection sampling, permits flexible step sizes, and runs with a deterministic runtime budget. Numerical experiments demonstrate that our approach converges rapidly to the target distribution, driven by interactions among multiple particles and by exploiting parallel computation.
Figures & tables
Figure 1: Illustration of the ideal proximal sampling (left) and our approximation (right). Heat flow and reverse dynamics are defined between πX and πY , but applied to intermediate ρ . Although these do not reach their targets in one step, the ideal version attains exponential convergence. Compared to the rejection sampling-based implementation of proximal samplers, our approach allows for larger step sizes (i.e., stronger convolution), which reduce the iterations to reach the target distribution.
Figure 2: Convergence of estimated KL divergence, averaged over 10 random seeds with shaded areas indicating variances. Our method (orange) outperforms both the proximal sampler with RGO (blue) and an ablated variant of our algorithm without particle interactions (green). It achieves the same accuracy as RGO in about 10× fewer iterations ( 100× faster when accounting for thinning).
Figure 3: One-dimensional marginals of πX along the third coordinate. The red curve is the ground-truth. Our method around 100 iterations (left, orange) already matches the RGO-based sampler at ∼ 1000 iterations (left, blue), and with only 200–300 iterations (right, orange) it closely aligns with the reference obtained from a sufficiently long run following Liang and Chen (2023b) (right, blue).
Figure 4: Empirical distributions at k=3,10,200 iterations with seed =0 on the two-tori domain. In-and-Out (top row) finds T1 , which overlaps with the initial standard Gaussian distribution, but fails to reach T2 . Our method (bottom row) generates some particles outside the domain but gradually drives particles toward T2 , demonstrating its ability to explore both components.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Parameters
Proximal Sampler with RGO
Initial distribution: y1/2(j)∼N(0,Id)
Step size: η=1/135
Number of independent chains: 100
Thinning: 10
Ours
Initial distribution: x0(i)∼N(0,Id)
Step size: h=1/10
Appendix
Table 1: Parameter setting for the experiment in Section 5.1 .
Method
Parameters
In-and-Out
Initial distribution: x0(i)∼N(0,Id)
Step size: h=1
Number of proposals for rejection sampling: 10000
(Particles are discarded if not accepted)
Number of independent chains: 1000
Ours
Initial distribution: x0(i)∼N(0,Id)
Appendix
Table 2: Parameter setting for the experiment in Section 5.2 .
Figure 5: Convergence of KL divergence for different step sizes h , with all other parameters set as Ours in Table 1 . Each curve shows the mean over 10 random seeds, with shaded areas indicating variances. Larger step sizes lead to faster convergence.
Figure 6: Comparison of hyperparameter tuning while keeping the number of particles used during the algorithm’s execution (i.e., M×N ). For readability, the first iteration has been omitted. While increasing the number of particles N used for approximating the target distribution stabilizes the KL divergence, changing both M and N results in very similar convergence patterns.
We study the problem of sampling from μ(dx)∝e−V(x)dx on Rd, where V is α-strongly convex and β-smooth, and write κ:=β/α. We design and analyze the Proximal Bouncy Particle Sampler (Proximal BPS), a new sampler that combines ideas from the proximal sampler and the bouncy particle sampler. From a warm start initialization with O(1) Rényi divergence w.r.t. μ, Proximal BPS returns a sample whose law is ε-close to μ in total variation distance using O(κd1/4polylog(1/ε)) gradient queries in expectation.
Fan Chen, Sinho Chewi, Jianfeng Lu +1
Department of Electrical Engineering and Computer Science, Massachusetts Institute of Technology. · Department of Statistics and Data Science, Yale University. · Department of Mathematics, Duke University. +1
We propose Decentralized Proximal Stochastic Gradient Langevin Dynamics (DE-PSGLD), a decentralized Markov chain Monte Carlo (MCMC) algorithm for sampling from a log-concave probability distribution constrained to a convex domain. Constraints are enforced through a shared proximal regularization based on the Moreau-Yosida envelope, enabling unconstrained updates while preserving consistency with the target constrained posterior. We establish non-asymptotic convergence guarantees in the 2-Wasserstein distance for both individual agent iterates and their network averages. Our analysis shows that DE-PSGLD converges to a regularized Gibbs distribution and quantifies the bias introduced by the proximal approximation. We evaluate DE-PSGLD for different sampling problems on synthetic and real datasets. As the first decentralized approach for constrained domains, our algorithm exhibits fast posterior concentration and high predictive accuracy.
Mohammad Rafiqul Islam, Lingjiong Zhu
Department of Mathematics, Florida State University, Tallahassee, Florida, United States of America
We study sampling from π(x)∝e−V(x) on Rd, where V is α-strongly convex and β-smooth with condition number κ=β/α. For this problem, the proximal sampler preserves the target exactly and recovers overdamped Langevin dynamics in the small-step limit, but its Gaussian resampling step discards directional information. Underdamped Langevin dynamics retains such information through momentum to accelerate exploration, motivating a proximal construction that combines target preservation with directional persistence. We introduce reflected proximal dynamics, which combines proximal sampling with an exact reflection of the auxiliary variable, conditional kinetic motion, and controlled refreshment. We develop two realizations: a flux-resampled underdamped half-turn and a persistent-momentum quarter-turn. From a warm start with constant χ2 divergence, both methods achieve total variation error ε using O(κd1/4) expected oracle queries. We further remove the supplied warm-start assumption by combining either sampler with existing posterior samplers in a proximal initialization stage with an optimized coarse step size. Accounting for both initialization and sampling yields O(κ6/7d1/4) expected oracle queries. In both regimes, the dependence on the target accuracy ε is polylog(1/ε), making reflected proximal dynamics a promising framework for accelerated high-accuracy log-concave sampling.