---
title: Parallel-Tempered SGHMC
url: https://www.emergentmind.com/topics/parallel-tempered-sghmc
type: topic
---

# Parallel-Tempered SGHMC

Parallel-Tempered Stochastic Gradient Hamiltonian Monte Carlo (Parallel-Tempered SGHMC), more precisely described as Parallel-Tempered Stochastic Gradient Nosé–Hoover Thermostat (PT-SGNHT), is a Markov Chain Monte Carlo (MCMC) method integrating parallel tempering with Nosé–Hoover (NH) thermostatted dynamics to efficiently address multimodal posterior sampling under stochastic gradients. By orchestrating a population of replicas at a geometrically-spaced ladder of temperatures and enabling periodic global state swaps, PT-SGHMC achieves reliable ergodic exploration of complex posterior landscapes, crucial for high-dimensional Bayesian inference and deep neural network learning with mini-batches [1812.01181].

## 1. Algorithmic Structure and Continuous-Time Dynamics

PT-SGHMC maintains $R$ independent replicas, each evolving under NH dynamics at distinct temperatures $T_1 < \ldots < T_R$. Each replica follows discretized NH evolution with parameters adapted to its temperature:
- The extended Hamiltonian for replica $i$ at inverse temperature $\beta_i = 1/T_i$ is
  $$
  H(\theta,p,\xi) = \beta_i U(\theta) + \frac{1}{2}p^\top M^{-1} p + \frac{1}{2} \xi^2 Q,
  $$
  where $U(\theta) = -\log \pi(\theta|\mathcal{D})$ and $\xi$ is the thermostat variable.
- The evolution of each replica is governed by
  $$
  d\theta = M^{-1} p\,dt, \\
  dp = -\beta_i \nabla U(\theta)\,dt - \xi\,p\,dt, \\
  d\xi = \frac{p^\top M^{-1} p - D}{Q}\,dt,
  $$
  with stochastic-gradient noise incorporated as an implicit Brownian term.
- In discretization:
  $$
  \theta_i \leftarrow \theta_i + \epsilon M^{-1} p_i, \\
  p_i \leftarrow p_i - \epsilon \frac{\nabla \tilde U(\theta_i)}{T_i} - \epsilon \xi_i p_i, \\
  \xi_i \leftarrow \xi_i + \epsilon \frac{p_i^\top M^{-1} p_i - D}{Q}.
  $$

Key algorithmic steps include alternation between independent NH updates and exchange moves, which swap full $(\theta, p, \xi)$ states between adjacent replicas, preserving detailed balance on the extended state space.

## 2. Exchange Moves and Parallel Tempering Protocol

In the parallel tempering framework, configuration swaps between replicas at neighboring temperatures enhance exploration by enabling low-temperature chains to inherit modes discovered by higher-temperature chains with higher mobility. The acceptance probability for exchanging states between replicas $(i, i+1)$ is determined using Barker's test:
- Compute $\Delta = (1/T_{i+1} - 1/T_i)[U(\theta_i) - U(\theta_{i+1})]$.
- Acceptance $\alpha = 1/(1 + \exp(-\Delta))$.
- Draw $u \sim \text{Uniform}(0,1)$; swap if $u < \alpha$.

Replica exchanges involve transferring the entire $(\theta, p, \xi)$ triplet, ensuring each configuration’s adaptive thermostat tracks the temperature history, which preserves ergodicity and reversibility.

## 3. Theoretical Guarantees and Multimodal Exploration

Each individual replica’s NH dynamics at temperature $T_i$ leaves the stationary distribution $\pi_i(\theta) \propto \exp(-U(\theta)/T_i)$ invariant and is ergodic. Swap proposals strictly maintain detailed balance for the joint product distribution $\prod_i \pi_i$. Thus, the overall chain remains reversible and ergodic, with the marginal at $T_1=1$ targeting $\pi(\theta|\mathcal{D})$.

High-temperature replicas facilitate movement across energy barriers, preventing multimodal trapping. Swapping propagates such diversified configurations to lower temperatures, a phenomenon unattainable through single-chain SGHMC or SGNHT under high-gradient noise.

## 4. Hyperparameterization and Computational Considerations

Hyperparameter selection is central for stability and efficiency:
- Temperature ladder is recommended geometric: $T_i = T_{\max}^{(i-1)/(R-1)}, \, T_{\max} \approx 10-100,\, R = 10-20$.
- Mass matrix $M$ as identity or diagonal (possibly estimated from past gradients).
- Thermostat inertia $Q$ ideally $O(D)$ for smooth adaptation.
- Step-size $\epsilon$ via dual-averaging pilot runs.
- Swap interval $L = 10-50$ iterations.
- Computational complexity per iteration is $O(R \cdot \text{cost}(\nabla \tilde U))$.
- Swap communication involves $O(D)$ floats between adjacent worker pairs.

These guidelines ensure scalability on large models and datasets. Packing replicas into a single $[R \times D]$ tensor facilitates batched GPU execution. Exchange of states is optimized using asynchronous point-to-point (MPI/NCCL) communication.

## 5. Empirical Performance on Synthetic and Real-World Problems

On a 1D 4-component Gaussian mixture (gradient noise $\sigma^2 = 0.25$), PT-SGNHT discovers all modes within approximately $10^3$ iterations, outperforming classic SGNHT/HMC which remain trapped. At $T=1$, effective sample size is tripled. In 2D, PT-SGNHT consistently finds all five modes, with mixing time halved relative to standard SGNHT.

For Bayesian neural networks (three-layer MLP, MNIST, batch size 100):
- PT achieves test log-likelihood $-0.093$ (compared to $-0.102$ for SGHMC and $-0.115$ for SGNHT).
- Accuracy increases by approximately $0.6\%$; predictive entropy decreases, indicating superior uncertainty quantification.

## 6. Implementation Details and Practical Tips

Optimizations include:
- Packing all $R$ replicas as a single tensor to exploit batched computation.
- Regularly estimating gradient noise variance $\sigma^2$ and adjusting $Q$ for robust thermostat adaptation.
- For noisy mini-batch acceptance under $U$, a Hermite-expansion correction $p_\mathcal{C}$ (with $k \leq 4$) mitigates swap bias.
- Non-blocking communication strategies (asynchronous MPI/NCCL) facilitate efficient local state swapping.
- Maintaining per-replica random seeds and checkpointing full $(\theta, p, \xi)$ states ensures reproducibility and convenient warm restarts.

## 7. Significance and Use Cases

Parallel-Tempered SGHMC/SGNHT addresses key bottlenecks in scalable Bayesian computation, especially for models with pronounced multimodality and data-partition-induced gradient noise. By integrating parallel tempering with adaptive NH thermostats, it enables principled, reliable posterior inference central to deep Bayesian learning for large-scale datasets, where conventional single-chain MCMC strategies are inadequate [1812.01181].

Source: https://www.emergentmind.com/topics/parallel-tempered-sghmc