Covariance-Controlled Adaptive Langevin Thermostat

• CCAdL thermostat is a stochastic sampling method that adaptively controls parameter-dependent covariance noise in Bayesian inference.
• It leverages instantaneous covariance estimation and dynamic friction modulation to ensure ergodic and robust sampling across high-dimensional landscapes.
• The modified CCAdL (mCCAdL) employs matrix-exponential integration for enhanced stability and faster convergence in large-scale, noisy simulations.

The Covariance-Controlled Adaptive Langevin (CCAdL) thermostat is an advanced stochastic diffusion-sampling algorithm designed for large-scale Bayesian posterior inference with parameter-dependent, state-varying gradient noise. Its formulation introduces exact, adaptive dissipation of covariance-structured noise inherent to stochastic gradient estimates and preserves the correct Gibbs invariant measure. CCAdL extends previously established thermostats by leveraging instantaneous covariance estimation coupled with dynamic friction modulation, providing robust and ergodic sampling for complex, high-dimensional posterior landscapes, especially in stochastic-gradient and quantum molecular simulations (Shang et al., 2015, Leimkuhler et al., 2015, Wei et al., 30 Dec 2025, Mouhat et al., 2017).

1. Background and Problem Setup

In Bayesian inference, the posterior density for a parameter vector $\theta\in\mathbb{R}^{N_d}$ given data $X=\{x_i\}_{i=1}^N$ is

$$\pi(\theta|X) \propto \pi(X|\theta)\,\pi(\theta) = \exp[-U(\theta)],$$

where $U(\theta) = -\log \pi(X|\theta) - \log \pi(\theta)$ is the posterior potential. Standard Langevin and Hamiltonian Monte Carlo (HMC) methods require access to full gradients $\nabla U(\theta)$, which is infeasible at scale. Stochastic gradient methods replace this with an unbiased estimator:

$$\nabla \tilde{U}(\theta) = \nabla U(\theta) + \xi(\theta), \qquad \xi(\theta)\sim\mathcal{N}(0,\Sigma(\theta)),$$

where the covariance $\Sigma(\theta)$ encodes mini-batch gradient noise and is generally parameter-dependent.

Conventional stochastic thermodynamic samplers introduce constant artificial noise or friction, assuming $\Sigma$ constant or known. However, parameter-dependent noise is generically required for accuracy and efficiency in large-scale or high-variance settings (e.g., modern ML, quantum simulations).

2. CCAdL Stochastic Dynamics and Stationary Law

The continuous-time CCAdL thermostat augments the classical underdamped Langevin SDE by:

1. An explicit friction drift $-\frac{1}{2}\Sigma(\theta)p\,dt$ that matches the Itô correction from $\theta$-dependent noise.
2. A Nosé–Hoover thermostat variable $\xi$ that adaptively adjusts friction to enforce target kinetic energy.

The SDE system with mass matrix $M=I$, unit temperature, and independent Wiener processes $W, W_A$ reads (Shang et al., 2015): $$\begin{aligned} &d\theta = p\,dt \ &dp = \left[-\nabla U(\theta) - \frac{1}{2}\Sigma(\theta)p - \xi\,p\right]\,dt + \sqrt{\Sigma(\theta)}\,dW + \sqrt{2A}\,dW_A \ &d\xi = \mu^{-1}\left[\frac{1}{N_d}p^T p - 1\right]\,dt \end{aligned}$$ The extended Gibbs invariant density is

$$\rho^*(\theta,p,\xi) \propto \exp\left[-\frac{1}{2}p^Tp - U(\theta)\right]\exp\left[-\frac{\mu}{2}(\xi-A)^2\right]$$

and is stationary under the corresponding Fokker–Planck operator, guaranteeing exact posterior marginals $\theta\sim\pi(\theta)\propto\exp[-U(\theta)]$. Ergodicity follows by Hörmander's condition due to strong coupling between $p$ and $\xi$ (Leimkuhler et al., 2015).

3. Discretization, Covariance Estimation, and Practical Integration

Evaluation of the parameter-dependent covariance $\Sigma(\theta)$ requires online estimation, which is typically achieved via an exponential moving average (EMA) of empirical mini-batch gradient covariances (Shang et al., 2015, Leimkuhler et al., 2015): $$I_t = (1-\kappa_t)I_{t-1} + \kappa_t V_t,\quad V_t = \frac{1}{n-1}\sum_{i=1}^n \left[g(\theta;x_{r_i}) - \bar g_t\right]\left[g(\theta;x_{r_i}) - \bar g_t\right]^T.$$ Discretization employs either a simple Euler-type step (original CCAdL) or higher-order symmetric splittings for improved order and stability (BADODAB, BAODCDOAB). Running statistics for covariance estimation enable CCAdL to immediately adjust friction and noise, mitigating bias that arises in fixed-covariance samplers such as SGLD or SGHMC. The BADODAB (symmetric SGNHT-S) splitting delivers second-order weak accuracy and fourth-order configurational “superconvergence” in the large friction limit (Leimkuhler et al., 2015).

4. Stability, Modified Integration (mCCAdL), and Algorithmic Advances

The original Euler discretization for the friction-covariance “C-part” entails step-size restrictions determined by the spectral radius of $\tilde\Sigma$. The modified CCAdL (mCCAdL) thermostat (Wei et al., 30 Dec 2025) replaces this with an efficient, stable matrix-exponential update using scaling and squaring alongside a truncated Taylor approximation: $$p_\text{C-step}(t) = \exp\left(t\,\tilde\Sigma\right)p(0),\quad \tilde\Sigma = -\frac{h}{2}\beta\Sigma(\theta)M^{-1}.$$ The full integrator employs a BAODCDOAB symmetric splitting, with each substep either solvable in closed form or by a fast expmv-type procedure: $$\begin{array}{rl} \text{(B)} & p \leftarrow p + (h/2)\tilde F(\theta) \ \text{(A)} & \theta \leftarrow \theta + (h/2)M^{-1}p\ \text{(O)} & p \leftarrow e^{-\xi h/2}p + \sqrt{\frac{A}{\beta}\frac{1-e^{-\xi h}}{\xi}}N(0,I)\ \text{(D)} & \xi \leftarrow \xi + (h/2)\mu^{-1}\left(p^T M^{-1}p - N_d/\beta\right)\ \text{(C)} & p \leftarrow \exp(h\tilde\Sigma(\theta))p \end{array}$$ Largest stable $h$ is increased by $10$–$20\times$; mCCAdL enables order-of-magnitude faster, more robust chains, especially for large $\Sigma$ (Wei et al., 30 Dec 2025).

5. Comparison to Related Stochastic Thermostats

A summary of distinctions:

Method	Covariance Handling	Friction Adaptivity	Stability/Stepper
SGLD	Scalar/constant	Fixed	Euler, highly limited
SGHMC	Estimated/constant	Fixed or tuned	Requires high friction
SGNHT	Constant	Adaptive (scalar)	1st/2nd order split
CCAdL	Full, parameter-wise	Adaptive (matrix)	Euler, BADODAB
mCCAdL	Full, parameter-wise	Adaptive (matrix)	Symmetric, matrix-exp

CCAdL uniquely adapts to variable $\Sigma(\theta)$, dissipates local noise, and avoids the bias introduced by methods not accounting for parametric covariance. In high dimensions, diagonal or low-rank approximations for $\Sigma$ maintain tractability (Shang et al., 2015, Wei et al., 30 Dec 2025).

6. Applications and Empirical Results

CCAdL has been empirically validated in:

• Gaussian mean–variance inference: Achieves lowest RMSE/mixing times versus SGHMC and SGNHT, which degrade for variable $\Sigma(\theta)$ or large $h$ (Shang et al., 2015).
• Bayesian logistic regression: On MNIST/CIFAR-10, CCAdL/mCCAdL converge $2$–$3\times$ faster in test log-likelihood per epoch compared with alternatives; stability for $h$ is $10$–$20\times$ greater for mCCAdL (Wei et al., 30 Dec 2025).
• Discriminative RBMs: Robust to large $h$ and small $A$; outperform SGLD/SGHMC.
• Quantum molecular simulations: Used as the PIOUD thermostat in quantum Monte Carlo path-integral Langevin dynamics, enabling efficient, fluctuation-dissipation–consistent nuclear dynamics for noisy QMC force fields (Mouhat et al., 2017).

7. Parameter Selection and Practical Recommendations

• Stepsize $h$: mCCAdL permits $h\gg$ (original CCAdL/SGNHT/SGHMC). Empirical values: $h\in[5\times10^{-4},1.2\times10^{-3}]$ for logistic models, up to $3\times10^{-2}$ for RBMs.
• Thermal mass $\mu$: Chosen comparable to $N_d$; larger $\mu$ yields slower adaptation, smaller $\mu$ greater variance.
• Noise strength $A$: Moderate $A\approx1$–$10$ best balances ergodicity and mixing.
• Covariance Estimation: Original methods utilize EMA; mCCAdL's matrix-exponential integration removes the need for moving-average windows, eliminating associated instability (Wei et al., 30 Dec 2025).

CCAdL and its modified integration are state-of-the-art for Bayesian sampling in noisy, high-dimensional, or stochastic-gradient settings. Their ability to adapt to local, state-dependent covariance structures consistently improves sampling accuracy, convergence, and robustness relative to all prior Langevin- and thermostat-based samplers (Shang et al., 2015, Leimkuhler et al., 2015, Wei et al., 30 Dec 2025, Mouhat et al., 2017).