Skip to content

Torchebm > Samplers > Hmc

Contents

Classes

  • HamiltonianMonteCarlo - Hamiltonian Monte Carlo sampler for efficient exploration of complex probability distributions.

API Reference

torchebm.samplers.hmc

Hamiltonian Monte Carlo Sampler Module.

This module provides a robust implementation of the Hamiltonian Monte Carlo (HMC) algorithm, a powerful Markov Chain Monte Carlo (MCMC) technique. By leveraging Hamiltonian dynamics, HMC efficiently explores complex, high-dimensional probability distributions, making it ideal for Bayesian inference and statistical modeling.

Key Features

  • Efficient sampling using Hamiltonian dynamics.
  • Customizable step sizes and leapfrog steps for fine-tuned performance.
  • Diagnostic tools to monitor convergence and sampling quality.
Module Components

Classes:

Name Description
HamiltonianMonteCarlo

Implements the Hamiltonian Monte Carlo sampler.
Usage Example

Sampling from a Gaussian Distribution

from torchebm.samplers.mcmc import HamiltonianMonteCarlo
from torchebm.energy_functions.energy_function import GaussianEnergy
import torch

# Define a 2D Gaussian energy function
energy_fn = GaussianEnergy(mean=torch.zeros(2), cov=torch.eye(2))

# Initialize HMC sampler
hmc = HamiltonianMonteCarlo(energy_fn, step_size=0.1, n_leapfrog_steps=10)

# Starting points for 10 chains
initial_state = torch.randn(10, 2)

# Run sampling
samples, diagnostics = hmc.sample_chain(initial_state, n_steps=100, return_diagnostics=True)
print(f"Samples: {samples.shape}")
print(f"Diagnostics: {diagnostics.keys()}")
Mathematical Foundations

Hamiltonian Dynamics in HMC

HMC combines statistical sampling with concepts from classical mechanics. It introduces an auxiliary momentum variable \( p \) and defines a Hamiltonian:

\[ H(q, p) = U(q) + K(p) \]
  • Potential Energy: \( U(q) = -\log \pi(q) \), where \( \pi(q) \) is the target distribution.
  • Kinetic Energy: \( K(p) = \frac{1}{2} p^T M^{-1} p \), with \( M \) as the mass matrix (often set to the identity matrix).

This formulation allows HMC to propose new states by simulating trajectories along the energy landscape.

Why Hamiltonian Dynamics?

  • Efficient Exploration: HMC uses gradient information to propose new states, allowing it to explore the state space more efficiently, especially in high-dimensional and complex distributions.
  • Reduced Correlation: By simulating Hamiltonian dynamics, HMC reduces the correlation between successive samples, leading to faster convergence to the target distribution.
  • High Acceptance Rate: The use of Hamiltonian dynamics and a Metropolis acceptance step ensures that proposed moves are accepted with high probability, provided the numerical integration is accurate.
Leapfrog Integration

Numerical Simulation of Dynamics

HMC approximates Hamiltonian trajectories using the leapfrog integrator, a symplectic method that preserves energy. The steps are:

  1. Momentum Half-Step: $$ p_{t + \frac{\epsilon}{2}} = p_t - \frac{\epsilon}{2} \nabla U(q_t) $$
  2. Position Full-Step: $$ q_{t + 1} = q_t + \epsilon M^{-1} p_{t + \frac{\epsilon}{2}} $$
  3. Momentum Half-Step: $$ p_{t + 1} = p_{t + \frac{\epsilon}{2}} - \frac{\epsilon}{2} \nabla U(q_{t + 1}) $$

Here, \( \epsilon \) is the step size, and the process is repeated for \( L \) leapfrog steps.

Acceptance Step

Metropolis-Hastings Correction

After proposing a new state \( (q_{t + 1}, p_{t + 1}) \), HMC applies an acceptance criterion to ensure detailed balance:

\[ \alpha = \min \left( 1, \exp \left( H(q_t, p_t) - H(q_{t + 1}, p_{t + 1}) \right) \right) \]

The proposal is accepted with probability \( \alpha \), correcting for numerical errors in the leapfrog integration.

Practical Considerations

Tuning Parameters

  • Step Size (\( \epsilon \)): Too large a step size can lead to unstable trajectories; too small reduces efficiency.
  • Number of Leapfrog Steps (\( L \)): Affects the distance traveled per proposal—balance exploration vs. computational cost.
  • Mass Matrix (\( M \)): Adjusting \( M \) can improve sampling in distributions with varying scales.

How to Diagnose Issues?

Use diagnostics to check: - Acceptance rates (ideal: 0.6–0.8). - Energy conservation (should be relatively stable). - Autocorrelation of samples (should decrease with lag).

Common Pitfalls

  • Low Acceptance Rate: If the acceptance rate is too low, it may indicate that the step size is too large or the number of leapfrog steps is too high. Try reducing the step size or decreasing the number of leapfrog steps.
  • High Correlation Between Samples: If samples are highly correlated, it may indicate that the step size is too small or the number of leapfrog steps is too few. Increase the step size or the number of leapfrog steps to improve exploration.
  • Divergence or NaN Values: Numerical instability or poor parameter choices can lead to divergent behavior or NaN values. Ensure that the energy function and its gradients are correctly implemented and that parameters are appropriately scaled.
Advanced Insights

Why HMC Outperforms Other MCMC Methods

HMC's use of gradients and dynamics reduces random-walk behavior, making it particularly effective for: - High-dimensional spaces. - Multimodal distributions (with proper tuning). - Models with strong correlations between variables.

Further Reading