Reflected dynamics speed up sampling from complex mathematical models

Reflected Proximal Dynamics for Accelerated High-Accuracy Log-Concave Sampling

Machine LearningArtificial Intelligence

Summary

Sampling from complicated mathematical shapes is important but can be slow. This paper studies a new way to sample that keeps important directional information to speed things up. The researchers created methods combining these ideas that use fewer steps to reach high accuracy. Their approach works well even without perfect starting information. This makes it a promising method for precise calculations involving certain types of smooth and curved mathematical functions.

What this means in practice

  • For machine learning engineers: Enable faster and more accurate probabilistic inference for models with smooth, convex objectives using reflected proximal sampling methods.
  • For statistical data scientists: Reduce computational cost in Bayesian posterior sampling where high precision is needed and the target distribution is log-concave and smooth.

Authors

Bo Yuan, Wei Guo, Yongxin Chen

Abstract

We study sampling from $π(x)\propto\mathrm{e}^{-V(x)}$ on $\mathbb{R}^d$, 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 \textit{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 $\varepsilon$ using $\widetilde O(\sqrtκd^{1/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 $\widetilde O(κ^{6/7}d^{1/4})$ expected oracle queries. In both regimes, the dependence on the target accuracy $\varepsilon$ is $\mathrm{polylog}(1/\varepsilon)$, making reflected proximal dynamics a promising framework for accelerated high-accuracy log-concave sampling.