Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Modeling the Universe
ASTR 596 - Fall 2025

Part 4: Advanced MCMC Methods — From Random Walks to Hamiltonian Flow

When Physics Meets Inference | Inferential Thinking Module 5 | ASTR 596

San Diego State University

“In mathematics you don’t understand things. You just get used to them.”
— John von Neumann

“Nature uses only the longest threads to weave her patterns, so that each small piece of her fabric reveals the organization of the entire tapestry.”
— Richard Feynman

Learning Outcomes

By the end of Part 4, you will be able to:

  1. Diagnose when random walk Metropolis-Hastings fails (strongly correlated posteriors) using visual and quantitative metrics

  2. Explain why gradients provide directional information that makes sampling exponentially more efficient

  3. Derive Hamiltonian Monte Carlo from Hamilton’s equations and recognize it as N-body simulation in parameter space

  4. Implement HMC by adapting your Project 2 leapfrog integrator to parameter space

  5. Compute gradients via finite differences and understand why automatic differentiation is superior

  6. Interpret HMC diagnostics (energy conservation, divergences, acceptance rates)

  7. Understand how NUTS adaptively tunes HMC to eliminate manual parameter tuning

  8. Appreciate why JAX (autodiff + JIT) revolutionized gradient-based inference

Prerequisites: Part 3 (Metropolis-Hastings, detailed balance, diagnostics), Module 3 Part 2 (Hamiltonian dynamics, phase space), Project 2 (leapfrog integrator), calculus (partial derivatives, chain rule, optimization)

Roadmap: Where We’re Going

This part builds on everything you’ve learned about MCMC, physics, and computation. Here’s the journey:

figure-md - Unknown Directive
:align: center

```mermaid
flowchart TD
    A[Metropolis-Hastings<br/>struggles with correlations] --> B[Random walks<br/>are inefficient]
    B --> C[Need directional<br/>information]
    C --> E[Use Gradients!<br/>∇ log p]

    E --> F[Hamiltonian<br/>dynamics]
    F --> G[HMC<br/>Leapfrog integrator]

    G --> H[NUTS<br/>Auto-tuning]
    H --> I[JAX +<br/>NumPyro]

    classDef problem fill:#BF616A,stroke:#2E3440,stroke-width:2px,color:#fff
    classDef insight fill:#A3BE8C,stroke:#2E3440,stroke-width:2px,color:#fff
    classDef solution fill:#5E81AC,stroke:#2E3440,stroke-width:2px,color:#fff
    classDef tools fill:#EBCB8B,stroke:#2E3440,stroke-width:2px,color:#2E3440

    class A problem
    class E insight
    class G solution
    class I tools
```

The narrative arc: We’ll start by diagnosing when Metropolis-Hastings struggles (correlated posteriors), discover why gradients are powerful (they point uphill!), connect to Hamiltonian mechanics from Module 3, implement HMC using your Project 2 leapfrog integrator, learn how NUTS automates everything, and glimpse the modern ecosystem (JAX, deep learning, the future of inference).

This isn’t just a better MCMC algorithm. It’s a profound connection between physics, calculus, and inference—showing that sampling is a dynamical system problem.

Part 1: The MCMC Revolution and Its Limits

Why MCMC Changed Everything

In Part 3, you learned Metropolis-Hastings—one of the most important algorithms in scientific computing. Let’s appreciate its impact:

The MCMC revolution (1953-present):

Why it’s revolutionary: MCMC turned Bayesian inference from a theoretical framework into a practical tool. Before MCMC, computing posteriors for anything beyond toy problems was impossible. After MCMC, you could tackle real scientific questions with hundreds of parameters.

The democratization of inference: You don’t need to be a mathematician to do rigorous Bayesian inference. Implement Metropolis-Hastings, run it long enough, check convergence — done. This is computational thinking at its finest.

The Problem: When Random Walks Get Lost

Metropolis-Hastings is universal — it can sample from any distribution. But universal doesn’t mean efficient. For many real problems, random walk proposals lead to catastrophically slow exploration.

Let me show you where it breaks down.

Consider a 2D posterior with strong correlation — a twisted Gaussian. This isn’t a pathological edge case; it’s what real posteriors often look like when parameters trade off (think Ωm\Omega_m vs hh in cosmology, or planet mass vs orbital inclination in exoplanet fitting).

The Twisted Gaussian: Parameters (θ1,θ2)(\theta_1, \theta_2) with posterior:

p(θ1,θ2∣D)∝exp⁡(−12(1−ρ2)[θ12−2ρθ1θ2+θ22])p(\theta_1, \theta_2 | D) \propto \exp\left(-\frac{1}{2(1-\rho^2)}\left[\theta_1^2 - 2\rho\theta_1\theta_2 + \theta_2^2\right]\right)
(2)

where ρ=0.95\rho = 0.95 is the correlation coefficient.

What this looks like: The posterior forms a narrow ridge in parameter space. High probability along θ2≈θ1\theta_2 \approx \theta_1, but very narrow perpendicular to the ridge. Moving along the ridge is easy (parameters change together). Moving perpendicular is hard (tightly constrained).

Why random walk Metropolis-Hastings struggles:

  1. Isotropic proposals: Standard M-H uses q(θ′∣θ)=N(θ,σ2I)q(\theta'|\theta) = \mathcal{N}(\theta, \sigma^2 I) — a circular Gaussian centered on current position. Proposals are equally likely in all directions.

  2. The step-size dilemma:

    • Small σ\sigma: Proposals along the ridge are accepted (good!) but make tiny progress. Proposals perpendicular are also small, so you never explore the width. → Slow mixing along ridge, never explore perpendicular.

    • Large σ\sigma: Proposals along the ridge make big jumps (good!) but proposals perpendicular step out of the ridge and get rejected (bad!). → Most proposals rejected, acceptance rate collapses.

  3. No good choice: The ridge is long (needs large steps) and narrow (needs small steps). Isotropic proposals can’t do both simultaneously.

The consequence: Effective Sample Size (ESS) is terrible. Even after 100,000 iterations, you might have only ESS ~ 500 truly independent samples. You’ve wasted 99.5% of your computation!

Quantitative diagnosis: Let’s look at what random walk MCMC produces for the twisted Gaussian with ρ=0.95\rho = 0.95:

Compare this to HMC (spoiler alert):

Same problem, two orders of magnitude difference in efficiency. This is why gradient-based methods matter.

figure-md - Unknown Directive
:align: center

<img src="figures/04-mod5-part4-fig1-mh-vs-hmc.png" alt="Comprehensive comparison of Metropolis-Hastings vs HMC" width="95%">

**Figure 1: Metropolis-Hastings vs Hamiltonian Monte Carlo on a Strongly Correlated Posterior.** *Top left:* Target distribution (2D Gaussian with correlation ρ=0.95) showing the narrow ridge structure that challenges random walk samplers. *Top right:* M-H produces scattered samples with low acceptance (44.8%), requiring 10,000 iterations to poorly cover the posterior. *Bottom left:* HMC follows smooth trajectories along the ridge with high acceptance (97.5%), needing only 1,000 samples for superior coverage. *Bottom right:* Marginal distributions show HMC (blue) accurately recovers the true distribution (dashed black) while M-H (red) shows poor coverage despite 10× more samples. HMC achieves ~60× better effective sample size for the same computational cost.
figure-md - Unknown Directive
:align: center

<img src="figures/04-fig2_trace_plots_comparison.png" alt="MCMC trace plots comparison" width="95%">

**Figure 2: Trace Plot Diagnostic — The "Hairy Caterpillar" Test.** *Left:* M-H trace plots show slow, correlated wandering with visible trends and poor mixing—this is what *bad* MCMC looks like. The chain gets stuck in regions and slowly drifts. *Right:* HMC trace plots show rapid, noisy exploration around the stationary distribution—this "hairy caterpillar" appearance indicates good mixing. Each parameter explores its full range quickly and independently. The dramatic difference in mixing is why HMC needs far fewer iterations to achieve convergence.
figure-md - Unknown Directive
:align: center

<img src="figures/04-fig3_autocorrelation_comparison.png" alt="Autocorrelation function comparison" width="95%">

**Figure 3: Autocorrelation Quantifies Sampling Efficiency.** *Left:* M-H autocorrelation decays slowly (τ ≈ 158 steps), meaning samples remain correlated for hundreds of iterations. This high autocorrelation time means you need ~17× more samples to get independent draws. *Right:* HMC autocorrelation drops to zero almost immediately (τ ≈ 9 steps), indicating samples are nearly independent after just a few iterations. This is the quantitative proof that HMC is dramatically more efficient—the autocorrelation time is the "price per independent sample" and HMC's price is 17× lower.

Part 2: Gradients — Why Direction Matters (🔴 Essential)

Before we dive into HMC, let’s revisit a concept from Calculus I that turns out to be shockingly important: gradients and optimization.

Flashback to Calculus I: How do you find the minimum of a function f(x)f(x)?

  1. Compute the derivative: f′(x)f'(x)

  2. Set it to zero: f′(x)=0f'(x) = 0

  3. Solve for x∗x^* (critical point)

  4. Check second derivative: f′′(x∗)>0f''(x^*) > 0 → local minimum

For multivariable functions f(θ1,θ2,…,θd)f(\theta_1, \theta_2, \ldots, \theta_d):

  1. Compute the gradient: ∇f=[∂f∂θ1,∂f∂θ2,…,∂f∂θd]\nabla f = \left[\frac{\partial f}{\partial \theta_1}, \frac{\partial f}{\partial \theta_2}, \ldots, \frac{\partial f}{\partial \theta_d}\right]

  2. Set it to zero: ∇f=0\nabla f = 0

  3. Solve for θ∗\theta^* (critical point)

The gradient is the multivariable generalization of the derivative. It points in the direction of steepest ascent — the direction where ff increases most rapidly.

Physical intuition: If f(θ)f(\theta) is a landscape’s elevation, ∇f(θ)\nabla f(\theta) is the slope vector at your location. It points uphill. The magnitude ∥∇f∥\|\nabla f\| tells you how steep the slope is.

Why Gradient Descent Is a Big Freaking Deal

Here’s the problem: For complicated functions, solving ∇f=0\nabla f = 0 analytically is impossible. Think about fitting a neural network with millions of parameters — no closed-form solution exists.

But here’s the trick: Even if you can’t solve ∇f=0\nabla f = 0, you can follow the gradient downhill to find a minimum. This is gradient descent:

θn+1=θn−ϵ∇f(θn)\theta_{n+1} = \theta_n - \epsilon \nabla f(\theta_n)
(3)

where ϵ\epsilon is the step size (learning rate).

Algorithm:

  1. Start at some θ0\theta_0 (initial guess)

  2. Compute gradient: g=∇f(θn)g = \nabla f(\theta_n)

  3. Take a step downhill: θn+1=θn−ϵg\theta_{n+1} = \theta_n - \epsilon g

  4. Repeat until ∥∇f∥≈0\|\nabla f\| \approx 0 (converged to a minimum)

Why this is profound:

This is how we train neural networks! Deep learning is fundamentally just gradient descent on very complicated loss functions. Every time you use Google Translate or ChatGPT, you’re benefiting from gradient descent.

From Optimization to Sampling: The Key Idea

Optimization: Find θ∗=arg⁡max⁡p(θ∣D)\theta^* = \arg\max p(\theta|D) (the single best parameter value)

Sampling: Generate samples θ(i)∼p(θ∣D)\theta^{(i)} \sim p(\theta|D) (explore the whole distribution)

Why sampling is harder: Optimization needs one point (the peak). Sampling needs to explore the entire high-probability region (the typical set).

But here’s the connection: The gradient ∇ln⁡p(θ∣D)\nabla \ln p(\theta|D) tells you:

Metropolis-Hastings ignores this information! It proposes randomly, not using the landscape’s structure.

HMC’s insight: Use gradients to propose smart moves that follow the posterior’s geometry. Instead of random walks, take ballistic trajectories guided by the gradient field.

Preview: HMC treats the negative log-posterior U(θ)=−ln⁡p(θ∣D)U(\theta) = -\ln p(\theta|D) as a potential energy and simulates dynamics in this landscape. Gradients become forces. This is where physics meets inference.

Computing Gradients: Two Approaches

To use gradients in HMC, we need to compute ∇θln⁡p(θ∣D)\nabla_\theta \ln p(\theta|D) for our posterior. Two main approaches:

Approach 1: Finite Differences (What You Already Know)

Recall from earlier in the semester: We approximated derivatives numerically using finite differences:

∂f∂θi≈f(θ+hei)−f(θ−hei)2h\frac{\partial f}{\partial \theta_i} \approx \frac{f(\theta + h \mathbf{e}_i) - f(\theta - h \mathbf{e}_i)}{2h}
(4)

where ei\mathbf{e}_i is the unit vector in the ii-th direction and h∼10−5h \sim 10^{-5} to 10-7 is a small step.

For a dd-dimensional parameter vector, compute all dd partial derivatives:

def grad_log_posterior_fd(theta, log_post_fn, h=1e-6):
    """Finite difference gradient of log-posterior."""
    d = len(theta)
    grad = np.zeros(d)
    
    for i in range(d):
        theta_plus = theta.copy()
        theta_minus = theta.copy()
        theta_plus[i] += h
        theta_minus[i] -= h
        
        grad[i] = (log_post_fn(theta_plus) - log_post_fn(theta_minus)) / (2*h)
    
    return grad

Cost: 2d2d function evaluations per gradient (one forward, one backward for each parameter).

Accuracy: Finite differences introduce truncation error:

Approach 2: Automatic Differentiation (The Modern Way)

Here’s the magic: What if you could compute exact gradients (up to machine precision) with the same cost as one function evaluation?

This is what automatic differentiation (autodiff) does. Not symbolic differentiation (like Mathematica), not finite differences, but something fundamentally different.

The idea: Every computer program is just a sequence of elementary operations (+, -, ×, ÷, exp, log, sin, etc.). Each operation has a known derivative. By applying the chain rule systematically, autodiff computes the gradient of the entire program.

Example: Compute gradient of f(θ1,θ2)=exp⁡(θ1θ2)+sin⁡(θ1)f(\theta_1, \theta_2) = \exp(\theta_1 \theta_2) + \sin(\theta_1)

Forward pass (compute function value):

v1 = θ₁ × θ₂
v2 = exp(v1)
v3 = sin(θ₁)
f = v2 + v3

Backward pass (compute gradient via chain rule):

∂f/∂v3 = 1
∂f/∂v2 = 1
∂f/∂v1 = ∂f/∂v2 × exp(v1) = exp(v1)
∂f/∂θ₂ = ∂f/∂v1 × θ₁ = θ₁ exp(θ₁θ₂)
∂f/∂θ₁ = ∂f/∂v1 × θ₂ + ∂f/∂v3 × cos(θ₁) = θ₂ exp(θ₁θ₂) + cos(θ₁)

Result: Exact gradient in approximately the same time as computing ff itself!

In JAX, this looks like:

import jax
import jax.numpy as jnp

def log_posterior(theta):
    # Your likelihood and prior computations here
    return log_like + log_prior

# Get gradient function automatically
grad_log_posterior = jax.grad(log_posterior)

# Use it
theta = jnp.array([1.0, 2.0])
gradient = grad_log_posterior(theta)  # Exact gradient, fast!

That’s it. JAX (and PyTorch, TensorFlow) automatically computes gradients for you. No finite differences, no hand-coded derivatives.

Why is autodiff transformative?

  1. Speed: Gradient computation is no longer the bottleneck. For d=1000d = 1000 parameters, autodiff is ~1000× faster than finite differences.

  2. Accuracy: No numerical error from finite differences. Gradients are exact (up to floating-point precision).

  3. Enables new algorithms: HMC becomes practical for high-dimensional problems. Neural networks with billions of parameters become trainable.

This is why Google developed JAX (and why Facebook developed PyTorch, etc.): Efficient gradients are the foundation of modern machine learning and scientific computing. Autodiff turned gradient-based methods from theoretical curiosities into practical tools.

What This Means for MCMC

Random walk Metropolis-Hastings: Proposes blindly, uses only p(θ∣D)p(\theta|D) (the posterior value).

Gradient-based MCMC (HMC): Proposes using ∇θln⁡p(θ∣D)\nabla_\theta \ln p(\theta|D) (the posterior gradient).

Information gain: A gradient tells you dd pieces of information (one per parameter dimension). Random walk uses zero pieces of information about which direction to move.

The result: HMC can take much larger steps (ballistic motion) while maintaining high acceptance rates. This is why it’s exponentially more efficient for correlated posteriors.

Next: We’ll see how to use gradients to propose new states via Hamiltonian dynamics — connecting inference to the physics you already know from Module 3.

Part 3: Hamiltonian Monte Carlo — Physics as Inference Engine (🔴 Essential)

The Profound Connection: Sampling Is a Dynamical System Problem

Here’s a question that seems philosophical but turns out to be deeply practical:

If we want to sample from a distribution p(θ∣D)p(\theta|D), what physical system has that distribution as its equilibrium state?

This is the key insight behind Hamiltonian Monte Carlo. Instead of hand-crafting proposals, let physics do it for you.

The answer: A system with Hamiltonian:

H(θ,p)=U(θ)+K(p)H(\theta, p) = U(\theta) + K(p)
(5)

where:

Why this works: The joint distribution over (θ,p)(\theta, p) is:

p(θ,p)∝exp⁡(−H(θ,p))=exp⁡(−U(θ))exp⁡(−K(p))p(\theta, p) \propto \exp(-H(\theta, p)) = \exp(-U(\theta)) \exp(-K(p))
(6)

Notice that this factors:

p(θ,p)=p(θ)⋅p(p)p(\theta, p) = p(\theta) \cdot p(p)
(7)

where:

The magic: If we sample from the joint distribution p(θ,p)p(\theta, p), then the marginal distribution of θ\theta is exactly p(θ∣D)p(\theta|D)!

Strategy:

  1. Introduce auxiliary momentum variables pp

  2. Sample from p(θ,p)p(\theta, p) via Hamiltonian dynamics

  3. Throw away the momenta pp, keep the positions θ\theta

  4. → Samples from p(θ∣D)p(\theta|D)!

Why introduce momentum at all? Because Hamiltonian dynamics lets us efficiently explore the joint space (θ,p)(\theta, p), and marginalizing over pp gives us samples from our target p(θ∣D)p(\theta|D).

Hamilton’s Equations in Parameter Space

In classical mechanics, Hamilton’s equations describe how positions and momenta evolve:

drdt=∂H∂p,dpdt=−∂H∂r\frac{d\mathbf{r}}{dt} = \frac{\partial H}{\partial \mathbf{p}}, \quad \frac{d\mathbf{p}}{dt} = -\frac{\partial H}{\partial \mathbf{r}}
(8)

For HMC, replace r→θ\mathbf{r} \to \theta and compute the derivatives:

dθdt=∂H∂p=M−1p\frac{d\theta}{dt} = \frac{\partial H}{\partial p} = M^{-1} p
(9)
dpdt=−∂H∂θ=−∂U∂θ=∇θln⁡p(θ∣D)\frac{dp}{dt} = -\frac{\partial H}{\partial \theta} = -\frac{\partial U}{\partial \theta} = \nabla_\theta \ln p(\theta|D)
(10)

Look at the second equation carefully: The force on momentum is the gradient of the log-posterior! This is how gradients enter HMC — as forces driving the dynamics.

Physical interpretation:

Compare to N-body:

HMC potential: U(θ)=−ln⁡p(θ∣D)U(\theta) = -\ln p(\theta|D)

Same structure, different interpretation!

The Leapfrog Integrator (Again!)

To simulate Hamiltonian dynamics, we need to discretize Hamilton’s equations. You already did this in Project 2!

The leapfrog integrator for N-body:

# Half-step velocity update
v = v + (dt/2) * acceleration(r)

# Full-step position update
r = r + dt * v

# Half-step velocity update
v = v + (dt/2) * acceleration(r)

The leapfrog integrator for HMC:

# Half-step momentum update
p = p + (epsilon/2) * grad_log_posterior(theta)

# Full-step position update
theta = theta + epsilon * p  # Assuming M = I

# Half-step momentum update  
p = p + (epsilon/2) * grad_log_posterior(theta)

It’s the same algorithm! Just different variable names:

N-body (Project 2)HMC (Now)
Position r\mathbf{r}Parameters θ\theta
Velocity v\mathbf{v}Momentum pp
Acceleration a=−∇U/m\mathbf{a} = -\nabla U / mGradient ∇θln⁡p(θ∣D)\nabla_\theta \ln p(\theta|D)
Timestep Δt\Delta tStep size ϵ\epsilon

Why leapfrog? Recall from Project 2:

  1. Symplectic: Preserves phase space volume (Liouville’s theorem)

  2. Time-reversible: Can run dynamics forward then backward to return to start

  3. Energy-conserving: Hamiltonian HH is conserved to O(ϵ2)O(\epsilon^2) error

These properties make leapfrog ideal for HMC:

The Complete HMC Algorithm

Now we can write down the full algorithm. Each HMC iteration has two phases:

Phase 1: Hamiltonian Dynamics (generate proposal)

  1. Sample fresh momentum: p(0)∼N(0,M)p^{(0)} \sim N(0, M)

  2. Store current state: (θold,pold)=(θ(i−1),p(0))(\theta_\text{old}, p_\text{old}) = (\theta^{(i-1)}, p^{(0)})

  3. Simulate dynamics for LL leapfrog steps with step size ϵ\epsilon:

    • Result: (θnew,pnew)(\theta_\text{new}, p_\text{new})

  4. Negate momentum (for detailed balance): pnew←−pnewp_\text{new} \gets -p_\text{new}

Phase 2: Metropolis Acceptance (accept or reject proposal)

  1. Compute Hamiltonian change:

    ΔH=H(θnew,pnew)−H(θold,pold)\Delta H = H(\theta_\text{new}, p_\text{new}) - H(\theta_\text{old}, p_\text{old})
    (11)

  2. Accept with probability:

    α=min⁡(1,exp⁡(−ΔH))\alpha = \min(1, \exp(-\Delta H))
    (12)

  3. If accepted: θ(i)=θnew\theta^{(i)} = \theta_\text{new}
    If rejected: θ(i)=θold\theta^{(i)} = \theta_\text{old}

Key insight: If the leapfrog integration were perfect (ϵ→0\epsilon \to 0), we’d have ΔH=0\Delta H = 0 → α=1\alpha = 1 → 100% acceptance! But finite ϵ\epsilon introduces discretization error, so α<1\alpha < 1.

Trade-off: Smaller ϵ\epsilon → better energy conservation → higher acceptance, but need more leapfrog steps to explore. Larger ϵ\epsilon → fewer steps needed but lower acceptance. Sweet spot: α≈60\alpha \approx 60-90%.

Why HMC Works: Energy Conservation and Acceptance

The acceptance probability depends on ΔH\Delta H:

α=min⁡(1,exp⁡(−ΔH))\alpha = \min(1, \exp(-\Delta H))
(13)

If ΔH\Delta H is small, acceptance is high. If ΔH\Delta H is large, acceptance is low.

What determines ΔH\Delta H? Discretization error in the leapfrog integrator. Recall from Project 2:

Design goals:

  1. Keep ϵ\epsilon small enough that ΔH≲1\Delta H \lesssim 1 (high acceptance)

  2. Make LL large enough that trajectory explores far (low autocorrelation)

  3. Balance: (Lϵ)(L \epsilon) should be ~1 autocorrelation length in parameter space

Typical values:

Why is HMC so much better? Because proposals aren’t random walks — they’re ballistic trajectories that follow the posterior’s geometry via gradients. You’re surfing the ridge instead of stumbling around it.

Part 4: HMC in Action — The Twisted Gaussian (🔴 Essential)

Seeing Is Believing: Visual Comparison

Let’s apply HMC to the twisted Gaussian from Part 1 and compare it to random walk Metropolis-Hastings. This will make the efficiency gain visceral.

Problem setup: 2D posterior with strong correlation (ρ=0.95\rho = 0.95)

p(θ1,θ2)∝exp⁡(−12(1−ρ2)[θ12−2ρθ1θ2+θ22])p(\theta_1, \theta_2) \propto \exp\left(-\frac{1}{2(1-\rho^2)}\left[\theta_1^2 - 2\rho\theta_1\theta_2 + \theta_2^2\right]\right)
(14)

Target: Generate 2,000 samples from this posterior.

Methods compared:

  1. Random walk Metropolis-Hastings (from Part 3)

  2. Hamiltonian Monte Carlo (what we just learned)

Results: Trace Plots

Random Walk M-H:

HMC:

Visual result: HMC trace plots look like white noise (good!), while M-H traces show clear structure (bad!).

[FIGURE: Side-by-side trace plots for θ1\theta_1 and θ2\theta_2, M-H vs HMC]

Results: Trajectory Visualization

Random walk M-H trajectories: Look like drunk particles bouncing around. Lots of rejected proposals perpendicular to ridge. Slow progress along ridge.

HMC trajectories: Smooth arcs following the ridge. Momentum carries parameters along high-probability regions. Far fewer rejected proposals.

[FIGURE: 2D posterior with overlaid trajectories showing M-H (random walk) vs HMC (ballistic)]

Results: Autocorrelation Functions

Random Walk M-H: ACF decays slowly, τ≈250\tau \approx 250 steps

HMC: ACF decays rapidly, τ≈4\tau \approx 4 steps

ESS for 2,000 samples:

[FIGURE: ACF plots for both methods]

Results: Corner Plot

Corner plots show the 2D posterior as a contour map.

Random Walk M-H: Even after 10,000 samples, posterior looks patchy (under-explored)

HMC: After 2,000 samples, posterior is smooth and well-resolved

[FIGURE: Corner plots comparing M-H (10k samples) vs HMC (2k samples)]

The conclusion: HMC is orders of magnitude more efficient for correlated posteriors. This isn’t theoretical — it’s a practical game-changer.

Tuning HMC: Step Size and Trajectory Length

HMC has two hyperparameters to tune: ϵ\epsilon (step size) and LL (trajectory length). Let’s understand how to set them.

Step Size ϵ\epsilon: The Energy Conservation Trade-off

Too small (ϵ→0\epsilon \to 0):

Too large (ϵ≫1\epsilon \gg 1):

Sweet spot (ϵ\epsilon tuned):

Adaptive tuning (used by NUTS, Stan, PyMC): During warmup, adjust ϵ\epsilon via dual averaging:

Trajectory Length LL: How Far to Simulate

Too short (L→0L \to 0):

Too long (L→∞L \to \infty):

Sweet spot (LL tuned):

The U-turn problem: For very long trajectories, Hamiltonian motion eventually curves back. This is inefficient — we want to stop right before the U-turn. This is what NUTS automates!

figure-md - Unknown Directive
:align: center

<img src="figures/04-mod5-part4-fig2-parameter-grid.png" alt="HMC parameter tuning grid" width="90%">

**Figure 4: HMC Parameter Tuning Landscape.** Grid search over step size (ε) and trajectory length (L) reveals the efficiency landscape. Gray shading shows Effective Sample Size (ESS)—darker is better. Black contours mark ESS values, while blue dashed lines show acceptance rate contours. The green thick contour marks the optimal region (ESS > 80% of maximum). The sweet spot lies around ε ≈ 0.1-0.2 and L ≈ 20-40, achieving both high acceptance (65-80%) and maximum ESS. This figure demonstrates why tuning matters: poorly chosen parameters can reduce efficiency by 10×. Modern samplers (NUTS, Stan) automatically find this sweet spot during warmup.

HMC Diagnostics: What to Check

Beyond the standard MCMC diagnostics (trace plots, R-hat, ACF), HMC has specific diagnostics:

1. Energy Conservation

Plot ΔH\Delta H over iterations:

plt.plot(delta_H_history)
plt.axhline(0, color='red', linestyle='--', label='Perfect conservation')
plt.ylabel('ΔH per trajectory')
plt.xlabel('Iteration')
plt.title('Energy Conservation Check')

What to look for:

If ΔH\Delta H is consistently large: Step size ϵ\epsilon is too big. Decrease it.

2. Acceptance Rate

Target: 60-90% acceptance

3. Divergences

Divergence: When ΔH\Delta H is extremely large (>> 1), indicating numerical instability. Often occurs in regions of high posterior curvature.

What to do:

  1. Decrease step size ϵ\epsilon

  2. Use better mass matrix MM (adapt to posterior covariance)

  3. Reparameterize model (remove pathological geometry)

Divergences are warnings: Your sampler is struggling in some regions of parameter space. Don’t ignore them!

4. Standard MCMC Diagnostics

Don’t forget the diagnostics from Part 3! HMC still needs:

Part 5: The No-U-Turn Sampler (NUTS) (🟡 Important)

The Problem: Fixed Trajectory Length Is Suboptimal

HMC requires choosing LL (number of leapfrog steps). But the optimal LL depends on where you are in parameter space!

In tight regions (high curvature, small typical set):

In broad regions (low curvature, large typical set):

The dilemma: There’s no single LL that’s optimal everywhere. You need adaptive trajectory length.

This is what the No-U-Turn Sampler (NUTS) solves.

The No-U-Turn Criterion: When to Stop

Key insight: Keep simulating Hamiltonian dynamics until the trajectory starts to double back (“U-turn”). Then stop.

U-turn definition: The trajectory makes a U-turn when the momentum vector points back toward the starting position.

Mathematically: Let (θstart,pstart)(\theta_\text{start}, p_\text{start}) be the initial state and (θend,pend)(\theta_\text{end}, p_\text{end}) be the current state after kk leapfrog steps.

U-turn occurs when:

(θend−θstart)⋅pend<0(\theta_\text{end} - \theta_\text{start}) \cdot p_\text{end} < 0
(15)

Intuition: The displacement vector (θend−θstart)(\theta_\text{end} - \theta_\text{start}) points from start to current position. Momentum pendp_\text{end} points in the current velocity direction. If their dot product is negative, you’re moving back toward the start → U-turn!

Visual analogy: Imagine a ball rolling on a curved surface:

  1. Starts at bottom of valley, rolls uphill (momentum carries it)

  2. Reaches maximum height, momentarily stops

  3. Starts rolling back down (U-turn!)

  4. → Stop simulating right before the U-turn (you’ve explored as far as possible in this direction)

figure-md - Unknown Directive
:align: center

<img src="figures/04-mod5-part4-fig3-uturn-criterion.png" alt="NUTS U-turn detection criterion" width="95%">

**Figure 5: NUTS U-turn Detection — Knowing When to Stop.** Three snapshots of a single HMC trajectory evolving through parameter space. *Left:* Before U-turn (step 13)—displacement from start (black arrow) and momentum (yellow arrow) point in similar directions (angle 34°), dot product is positive, no U-turn detected. *Middle:* U-turn detected (step 27)—angle reaches 91°, dot product becomes negative (-0.04), NUTS would stop here. *Right:* After U-turn (step 47)—trajectory has doubled back (angle 133°), continuing wastes computation. The criterion $(θ_t - θ_0) · p_t < 0$ provides a simple, effective test for when the trajectory stops being productive. This automatic stopping is what makes NUTS "no-U-turn"—it explores as far as useful, then stops before wasting compute on redundant sampling.

The NUTS Algorithm (Conceptual Overview)

NUTS adaptively grows the trajectory until a U-turn is detected. Here’s the high-level idea:

NUTS procedure:

  1. Initialize: Sample momentum p∼N(0,M)p \sim N(0, M), set trajectory to current state (θ,p)(\theta, p)

  2. Build trajectory recursively:

    • Simulate forward in time (positive direction): Run leapfrog to extend trajectory

    • Simulate backward in time (negative direction): Run leapfrog in reverse

    • Double trajectory length each iteration: L=1,2,4,8,16,…L = 1, 2, 4, 8, 16, \ldots

    • Stop when U-turn detected

  3. Select proposal: Randomly choose a state from the trajectory (weighted by posterior probability)

  4. Metropolis acceptance: Accept/reject based on average ΔH\Delta H over trajectory

Key features:

The result: Fully automatic HMC with minimal tuning. This is why NUTS is the default in Stan and PyMC.

NUTS Advantages Over Vanilla HMC

AspectHMC (fixed LL)NUTS (adaptive)
User tuningMust choose LL manuallyAutomatic LL adaptation
EfficiencySuboptimal LL in many regionsNear-optimal LL everywhere
RobustnessSensitive to LL choiceWorks across diverse posteriors
ComplexitySimple algorithmMore complex (recursive doubling)

When to use NUTS:

When vanilla HMC is okay:

For Project 4: Start with vanilla HMC (easier to understand and debug). Extensions could explore NUTS via NumPyro.

The Complete Picture: Comparing All Three Methods

Now that you’ve seen the full progression from random walks to adaptive Hamiltonian flow, let’s compare all three methods side-by-side:

AspectMetropolis-HastingsHMC (vanilla)NUTS
Core ideaRandom walk proposalsHamiltonian dynamicsAdaptive HMC
Uses gradients?❌ No✅ Yes✅ Yes
Proposal mechanismθ′∼N(θ,σ2I)\theta' \sim N(\theta, \sigma^2 I)Leapfrog integrationAdaptive leapfrog with doubling
Tuning parametersStep size σ\sigmaStep size ϵ\epsilon, trajectory length LLStep size ϵ\epsilon (auto-tuned)
Typical acceptance23-50%60-90%65-95%
AutocorrelationHigh (τ∼\tau \sim 100-1000)Low (τ∼\tau \sim 3-10)Very low (τ∼\tau \sim 2-5)
ESS efficiency~1-5% of samples~10-30% of samples~20-50% of samples
Cost per iterationVery cheap (1 eval)Moderate (LL gradient evals)Moderate-expensive (adaptive LL)
Best forSimple, uncorrelated posteriorsCorrelated posteriors (d < 100)General purpose, production
Worst forStrongly correlated, high-dWhen gradients unavailableLow-d (overhead not worth it)
Requires derivatives?❌ No✅ Yes (finite diff or autodiff)✅ Yes (autodiff recommended)
RobustnessWorks everywhere (slowly)Sensitive to ϵ\epsilon, LLVery robust
Implementation complexitySimple (~20 lines)Moderate (~50 lines)Complex (~200 lines)
When posteriors multi-modalCan work (eventually)Can get stuckCan get stuck
Scales to high-d?❌ No (d > 20 struggles)⚠️ Moderate (d ~ 100)✅ Yes (d ~ 1000s with autodiff)
Professional toolsemcee (astronomy)Custom implementationsStan, PyMC, NumPyro

Key insights from the comparison:

  1. There’s no free lunch: HMC/NUTS require gradients and careful energy conservation, but reward you with exponentially better efficiency for hard problems.

  2. Cost vs efficiency trade-off:

    • M-H: Cheap iterations, but need many → Total cost high for correlated posteriors

    • HMC/NUTS: Expensive iterations (gradient evals), but need far fewer → Total cost lower

  3. When M-H is actually fine: If your posterior is simple (uncorrelated, low-dimensional), M-H works perfectly well. Don’t over-engineer!

  4. Why NUTS won: By eliminating manual tuning of LL, NUTS made gradient-based MCMC accessible to non-experts. This is why it’s the default in Stan and PyMC.

  5. The autodiff revolution: NUTS + JAX autodiff scales to thousands of parameters. This combination powers modern Bayesian deep learning and inference in complex simulators.

Practical decision tree:

Is your posterior correlated or high-dimensional?
│
├─ No (simple, uncorrelated, d < 10)
│  └─→ Use Metropolis-Hastings
│     └─→ Simple, fast, good enough
│
└─ Yes (correlated or d > 10)
   │
   ├─ Can you compute gradients?
   │  │
   │  ├─ No (black-box likelihood, discrete parameters)
   │  │  └─→ Use ensemble MCMC (emcee)
   │  │     └─→ No gradients, but uses multiple walkers
   │  │
   │  └─ Yes
   │     │
   │     ├─ Learning/prototyping?
   │     │  └─→ Implement vanilla HMC
   │     │     └─→ Understand the algorithm deeply
   │     │
   │     └─ Production/research?
   │        └─→ Use NUTS (Stan/PyMC/NumPyro)
   │           └─→ Automatic tuning, battle-tested
   │
   └─ Is posterior multimodal?
      └─→ All methods struggle!
         └─→ Try: Tempered MCMC, nested sampling, or VI

The bottom line: For Project 4 (cosmology with correlated Ωm\Omega_m and hh), HMC will dramatically outperform M-H. You’ll see this firsthand!

Beyond NUTS: The Broader MCMC Landscape

NUTS is state-of-the-art for general-purpose MCMC, but it’s not the end of the story. Here’s a glimpse of other advanced methods:

Riemannian Manifold HMC: Adapts the mass matrix MM locally based on posterior curvature (Hessian). Especially powerful for posteriors with highly varying scales.

Ensemble MCMC (affine-invariant, emcee): Multiple walkers explore together. No gradients needed! Good for multimodal posteriors. Used widely in astrophysics.

Variational Inference (VI): Approximate posterior by optimization (not sampling). Much faster than MCMC but less accurate. Trade-off: speed vs exactness.

Neural network-assisted MCMC: Learn proposal distributions via deep learning. State-of-the-art for extremely high dimensions (thousands of parameters).

When to use what:

The frontier: Combining MCMC with machine learning (learning optimal proposals, amortized inference, simulation-based inference). This is where Module 6 leads!

Part 6: JAX and the Future of Gradient-Based Inference (🟡 Important)

Why JAX Matters

You’ve seen that HMC needs gradients. Computing gradients efficiently is the bottleneck for scaling to high dimensions. This is where JAX transforms the game.

JAX is Google’s library for high-performance numerical computing. It brings three superpowers:

  1. Automatic Differentiation (Autodiff): Compute exact gradients for any Python function via jax.grad()

  2. Just-In-Time (JIT) Compilation: Compile Python to optimized machine code via jax.jit()

  3. Hardware Acceleration: Same code runs on CPU, GPU, TPU with no changes

Why Google developed JAX: Modern machine learning requires efficient gradients for training neural networks with billions of parameters. JAX enables this at scale.

Why you should care: JAX makes gradient-based inference (HMC, VI, ML) practical for high-dimensional problems. In Module 6, you’ll use JAX for physics-informed neural networks and differentiable simulations.

JAX Autodiff: A Quick Tour

The power of autodiff: Recall from Part 2 that automatic differentiation computes exact gradients with the cost of one function evaluation (not 2d2d like finite differences).

In JAX, getting gradients is trivial:

import jax
import jax.numpy as jnp

# Define a function
def f(x):
    return jnp.sum(x**2)

# Get its gradient (automatic!)
grad_f = jax.grad(f)

# Use it
x = jnp.array([1.0, 2.0, 3.0])
gradient = grad_f(x)  # Returns [2.0, 4.0, 6.0]

That’s it. JAX computes ∇f\nabla f automatically by tracing the computation graph and applying the chain rule.

For HMC, this means:

# Your log-posterior function
def log_posterior(theta):
    log_like = # ... your likelihood ...
    log_prior = # ... your prior ...
    return log_like + log_prior

# Get gradient function (one line!)
grad_log_posterior = jax.grad(log_posterior)

# Use in HMC
theta = jnp.array([0.3, 0.7])  # Current parameters
grad = grad_log_posterior(theta)  # Gradient for HMC

Speed comparison (d = 100 parameters):

For high-dimensional problems, this is the difference between hours and minutes.

JAX JIT: Make Code Fast

Just-In-Time compilation: JAX can compile Python functions to optimized machine code (like C++ or Fortran) automatically.

Example:

@jax.jit  # Compile this function!
def leapfrog_step(theta, p, epsilon, grad_fn):
    p = p + (epsilon/2) * grad_fn(theta)
    theta = theta + epsilon * p
    p = p + (epsilon/2) * grad_fn(theta)
    return theta, p

# First call: JAX compiles the function (takes a moment)
theta, p = leapfrog_step(theta, p, epsilon, grad_log_post)

# Subsequent calls: Run compiled code (10-100× faster!)
for _ in range(1000):
    theta, p = leapfrog_step(theta, p, epsilon, grad_log_post)

Speedup: 10-100× for numerical code (loops, array operations). The more complex your code, the bigger the speedup.

GPUs for free: Change jax.jit to run on GPU? Nothing! JAX handles device placement automatically. Your HMC code runs on GPU without modification.

The JAX Ecosystem for Inference

NumPyro: Probabilistic programming in JAX. Uses NUTS under the hood.

import numpyro
from numpyro.infer import MCMC, NUTS

def model(x, y):
    # Define your Bayesian model
    slope = numpyro.sample('slope', dist.Normal(0, 10))
    intercept = numpyro.sample('intercept', dist.Normal(0, 10))
    sigma = numpyro.sample('sigma', dist.HalfNormal(1))
    
    mu = intercept + slope * x
    numpyro.sample('y', dist.Normal(mu, sigma), obs=y)

# Run NUTS (fully automatic!)
nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_warmup=1000, num_samples=2000)
mcmc.run(jax.random.PRNGKey(0), x=x_data, y=y_data)

# Get results
samples = mcmc.get_samples()

BlackJAX: Lower-level MCMC library (like your from-scratch implementation, but production-ready). Modular design—mix and match samplers, diagnostics, adaptation strategies.

Why learn JAX now: Module 6 uses JAX extensively for differentiable physics simulations and physics-informed neural networks. Getting comfortable with JAX syntax now pays dividends later.

The Big Picture: Where Inference Is Heading

The convergence of fields: Gradient-based inference connects:

Modern inference workflow:

  1. Write model in JAX (probabilistic program)

  2. Automatic gradients via jax.grad()

  3. JIT-compile for speed

  4. Run NUTS on GPU (handles tuning automatically)

  5. Diagnose and visualize (arviz, corner.py)

This is how professionals do inference today. You’re learning the cutting edge.

Looking ahead:

The glass-box journey: You built MCMC from scratch (Part 3), understood the physics (Part 4), implemented HMC, and now see the professional tools (JAX, NumPyro). You’re not a black-box user — you’re a computational scientist who understands the foundations.

Part 7: Synthesis and Looking Forward

What We’ve Learned: The Journey from Random Walks to Hamiltonian Flow

Part 3 (Metropolis-Hastings):

Part 4 (HMC):

The connections:

ConceptModule 1Module 3Module 5 (Part 3)Module 5 (Part 4)
What we sampleRandom variablesMicrostatesParameters (random walk)Parameters (Hamiltonian)
Target distributionPopulationBoltzmannPosterior p(θ|D)Posterior p(θ|D)
How we sampleIID drawsMolecular dynamicsMarkov chain (M-H)Markov chain (HMC)
Key propertyLaw of Large NumbersErgodicityDetailed balanceHamiltonian conservation
DynamicsStaticNewton’s lawsStochastic jumpsLeapfrog integration
EfficiencyPerfect (independent)Limited by collision rateLimited by correlationHigh (gradient-guided)

The profound unity: Whether you’re:

You’re using variations of the same mathematical structure: stochastic or deterministic processes converging to equilibrium distributions, with time averages equaling ensemble averages.

This is why physics is powerful. The same principles work everywhere — from atoms to galaxies to probability distributions.

Where This Leads: Your Path Forward

Immediate: Project 4

You’ll apply HMC to real cosmology data:

Skills you’ll use:

Near future: Module 6

Physics-informed learning:

Long term: Your research

These tools are universally applicable:

The glass-box philosophy in action: You didn’t just learn to use Stan or PyMC. You:

→ You’re not a black-box user. You’re a computational scientist.

Self-Assessment Rubric

Before moving to Project 4, assess your mastery of Part 4. Be honest — identifying gaps now makes the project more productive.

Level 1: Conceptual Understanding

Can you explain the ideas to someone else?

Level 2: Physics Connection

Can you see the unity across modules?

Level 3: Mathematical Foundations

Can you derive key results?

Level 4: Implementation Skills

Can you code HMC from scratch?

Level 5: Diagnostic Expertise

Can you tell if your HMC sampler is working?

Level 6: Tool Awareness

Do you understand the modern ecosystem?

Level 7: Connections and Transfer

Can you see how HMC fits into the bigger picture?

Reflection: Which levels feel solid? Which need more work? Use this to guide your Project 4 strategy and what to review.

Conceptual Checkpoints (Throughout)

Looking Ahead: Project 4 and the Universe

You now have the tools to measure the contents of the universe. Seriously.

Project 4 preview: Type Ia supernovae cosmology

This is real science. The same analysis won the 2011 Nobel Prize in Physics for discovering cosmic acceleration (dark energy). You’re not just practicing methods—you’re doing frontier research with public data.

What you’ll experience:

Skills you’ll cement:

Beyond this course: These methods generalize to any inference problem — exoplanets, stellar evolution, black hole mergers, climate models, drug discovery, economics, you name it. Gradient-based MCMC is a universal tool.

Module 6 awaits: Physics-informed neural networks, differentiable simulations, combining data and theory. JAX becomes central. The journey continues from inference to learning.

Now go forth and sample the universe. 🌌

Further Resources (Optional)

For deeper understanding:

For implementation:

For inspiration:

Historical:

Next: Project 4 — Measuring dark energy with HMC 🚀

Bayesian Statistics Inference
Part 3: Markov Chain Monte Carlo — From Theory to Practice
Module 2 Computing The Universe
Overview: The JAX Revolution in Scientific Computing