Accelerated sampling using SamAdams variable timesteps and position-adaptive Langevin dynamics

Technical Summary: Accelerated Sampling using SamAdams Variable Timesteps and Position-Adaptive Langevin Dynamics

Problem Statement
The paper addresses the challenge of efficiently sampling from high-dimensional probability distributions defined by underdamped Langevin dynamics, particularly in landscapes with mixed stiffness (regions of high curvature alongside flat regions). Standard fixed-stepsize integrators face a fundamental trade-off: to maintain stability in stiff regions, the stepsize must be small enough to resolve the highest frequency modes, which leads to inefficient sampling in softer regions. Conversely, adaptive timestepping methods that contract steps only where necessary struggle to preserve the canonical invariant measure without complex reweighting or divergence corrections. Furthermore, isotropic friction matrices (damping all directions equally) often fail to optimize mixing, as they cannot independently tune damping along the force direction versus directions tangent to equipotential contours.

Methodology
The authors propose a combined scheme, SA-PAL, which integrates two complementary devices into a single palindromic integrator called AZBOBZA. This integrator requires only one force evaluation per iteration, maintaining the computational cost of standard second-order Langevin integrators.

1. SamAdams Adaptive Timestepping:

   • This device uses a "stiffness monitor" (either force magnitude $|\nabla U|$ or an arc-length metric involving velocity) to sense local dynamics.
   • A relaxation ODE governs a sensor variable $Z$ that tracks the monitor, smoothing out noise.
   • The effective stepsize $h_{eff} = h \cdot \eta$ is determined by a Sundman kernel $\psi(Z)$, which contracts the step in stiff regions and relaxes it elsewhere.
   • To recover the canonical distribution, Monte Carlo averages are computed using a time-reweighting procedure based on the accumulated scale factors $\eta$.

2. Position-Adaptive Langevin (PAL) Dynamics:

   • The friction matrix $\Gamma(x)$ is defined as a normalized rank-one-plus-scalar form: $\Gamma(x) = \alpha I + \beta \hat{f}(x)\hat{f}(x)^T$, where $\hat{f}(x)$ is the normalized force direction.
   • This structure ensures that friction is anisotropic: it applies a constant damping $\alpha$ in directions transverse to the force and $\alpha + \beta$ parallel to the force.
   • Crucially, this specific form allows for an exact $O(d)$ solution of the Ornstein–Uhlenbeck (OU) momentum update without requiring matrix decomposition or divergence corrections. The canonical distribution remains the exact invariant measure for any smooth, positive-definite $\Gamma(x)$ of this form.

3. The AZBOBZA Integrator:

   • The scheme combines the position update (A), momentum kick (B), exact OU step (O), and the adaptive sensor update (Z) into the sequence AZBOBZA.
   • This ordering ensures that the force $F = -\nabla U(x)$ is evaluated only once per step, serving both the momentum kicks and the friction direction calculation.
   • The integrator is weakly second-order accurate.

Key Contributions

• Algorithmic Efficiency: The authors demonstrate that combining variable timesteps with position-dependent friction can be achieved with no additional computational overhead compared to standard integrators, leveraging the rank-one structure of the friction tensor for exact updates.
• Preservation of Invariance: Unlike many preconditioned Langevin methods that require complex divergence corrections to preserve the target measure, the PAL friction form guarantees the canonical distribution is the exact invariant measure without such corrections.
• Ablation and Tuning Analysis: The paper provides a rigorous dissection of the contributions of the adaptive timestep versus the adaptive friction, showing that while the timestep handles stability and global exploration, the adaptive friction significantly improves mixing and accuracy in complex landscapes by decoupling damping rates.

Results
The method was tested on four benchmark problems:

1. Rosenbrock Function: SA-PAL improved mixing rates by 1.5–3 times compared to fixed-step integrators. The arc-length monitor was particularly effective in handling the stiff transverse modes while maintaining fast valley exploration.
2. Thin Entropic Channel: In this regime, the fixed-step integrators failed to cross the narrow neck efficiently. SA-PAL achieved efficiency gains of more than an order of magnitude, crossing the neck approximately 16 times more often than the best-tuned fixed-step baseline.
3. Mueller-Brown Potential: SA-PAL achieved a 2.8x faster mixing of positions compared to the best-mixing isotropic friction baseline, while simultaneously reducing potential-energy autocorrelation times. The adaptive friction allowed the system to damp energy-carrying modes without overdamping the slow inter-well transitions.
4. Bayesian Linear Regression (Horseshoe Prior): Applied to a high-dimensional problem with a "funnel" geometry, SA-PAL reduced the integrated autocorrelation time (IACT) for the aggregate coefficient observable by factors of 1.88 to 2.70 compared to accuracy-matched baselines, and by up to 52x compared to fixed-step BAOAB.

Significance and Claims
The paper claims that SA-PAL achieves a more favorable balance between exploration speed and equilibrium accuracy than isotropic-friction baselines, particularly in landscapes with curved low-energy structures or entropic bottlenecks. The authors emphasize that the primary advantage of the PAL component is the directional decoupling of damping rates, which prevents the overdamping of slow collective motions while rapidly thermalizing energy-changing directions.

The authors note that while the method offers substantial gains, the tuning of parameters (specifically the friction anisotropy $\beta$ and the sensor relaxation rate $\theta$) is non-trivial. They highlight that comparisons must be made at matched accuracy (bias), as adaptive schemes can appear faster but biased if the nominal stepsize is too aggressive. The paper concludes that SA-PAL is a robust, low-overhead method for accelerating sampling in complex, stiff potentials, though principled, problem-independent tuning protocols remain an open area for future work.