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.

Adjoint-Differentiated Laplace HMC

This notebook demonstrates blackjax.laplace_hmc, a sampler that integrates out latent Gaussian variables via the Laplace approximation and runs HMC over the marginal hyperparameter posterior.

Core idea. In a hierarchical model

phi   ~ p(phi)             # hyperparameters  (low-dimensional)
theta ~ N(0, K(phi))       # latent variables  (high-dimensional)
y     ~ p(y | theta, phi)  # likelihood

sampling the joint (θ,ϕ)(\theta, \phi) posterior is hard because the geometry changes dramatically with ϕ\phi (“funnel” geometry). Laplace-HMC instead runs HMC over ϕ\phi alone, using the Laplace approximation to analytically handle θ\theta. The mode θ∗(ϕ)=argmaxθlog⁡p(θ∣ϕ,y)\theta^*(\phi) = \text{argmax}_\theta \log p(\theta | \phi, y) is found via L-BFGS, and the log-marginal gradient uses the implicit function theorem — no gradient unrolling through the optimizer.

This notebook demonstrates:

  1. Efficiency on Neal’s Funnel — where Laplace HMC excels by marginalizing away the difficult geometry.

  2. GP Regression and Accuracy — where the Laplace approximation is exact.

  3. Comparison of Sampler Variants — benchmarking fixed/dynamic and standard/multinomial variants.

Section 1: Neal’s Funnel (Speed and Efficiency)

Neal’s Funnel is a classic example of a “difficult” geometry for MCMC. The joint distribution p(θ,v)p(\theta, v) is:

v∼N(0,32)θi∼N(0,exp⁡(v)2)for i=1…n\begin{aligned} v &\sim N(0, 3^2) \\ \theta_i &\sim N(0, \exp(v)^2) \quad \text{for } i=1 \dots n \end{aligned}

In the joint space, as vv becomes small, the conditional distribution of θ\theta becomes very narrow, forcing HMC to take tiny steps. Laplace HMC marginalizes out θ\theta analytically (the approximation is exact here!), leaving a simple Gaussian for vv.

Laplace HMC vs NUTS

Running Laplace-HMC window adaptation on Funnel...
Running Laplace-HMC sampling on Funnel...
Running NUTS on Funnel...
Laplace-HMC time: 79.69s
NUTS time:        18.92s
<Figure size 960x480 with 1 Axes>

Section 2: Gaussian Process Regression (Laplace is Exact)

In a GP regression model with Gaussian likelihood the Laplace approximation is exact: the conditional posterior p(θ∣ϕ,y)p(\theta \mid \phi, y) is itself Gaussian, so the marginal log⁡p(ϕ∣y)\log p(\phi \mid y) is computed without any approximation error. Laplace-HMC samples the two kernel hyperparameters ϕ=(log⁡ℓ,log⁡σf)\phi = (\log\ell, \log\sigma_f); the nn-dimensional function values θ\theta are recovered afterwards via laplace.sample_theta.

ϕj∼N(0,1),θ∣ϕ∼N(0,Kϕ),yi∣θi∼N(θi,σn2)\phi_j \sim \mathcal{N}(0, 1), \qquad \theta \mid \phi \sim \mathcal{N}(0, K_\phi), \qquad y_i \mid \theta_i \sim \mathcal{N}(\theta_i, \sigma_n^2)

where KϕK_\phi is the RBF (squared-exponential) kernel. NUTS must explore the full (n+2)(n+2)-dimensional space (θ,ϕ)(\theta, \phi) whose geometry is funnel-shaped; Laplace-HMC only explores the 2-dimensional ϕ\phi space.

Sampling full (32-D) joint with NUTS ...
Sampling full joint with NUTS done in 68.8s
Sampling GP hyperparameters with Laplace-HMC ...
Sampling GP hyperparameters with Laplace-HMC done in 74.5s (8x400 samples)
Source
<Figure size 1680x600 with 2 Axes>
Time / sample — Laplace-HMC: 23 ms  NUTS: 21 ms  (0.9x speedup)
log length-scale:  Laplace 0.23  NUTS 0.23
log amplitude:     Laplace -0.08  NUTS -0.07

Section 3: Comparing All Four Laplace-HMC Variants

The four top-level aliases form a 2×2 matrix along two independent axes:

Endpoint + M-HMultinomial trajectory
Fixed num_integration_stepsblackjax.laplace_hmcblackjax.laplace_mhmc
Random step count per transitionblackjax.laplace_dhmcblackjax.laplace_dmhmc

Axis 1 — fixed vs. dynamic step count. laplace_hmc and laplace_mhmc take a fixed num_integration_steps and plug directly into window_adaptation. laplace_dhmc and laplace_dmhmc draw the number of leapfrog steps uniformly at random each transition, removing periodic-orbit sensitivity; they require a rng_key argument at .init() and do not support window_adaptation.

Axis 2 — standard M-H vs. multinomial proposal. Standard variants (hmc/dhmc) propose the trajectory endpoint and apply an M-H accept/reject step; info.acceptance_rate is the usual Metropolis acceptance probability (target ≈ 0.65). Multinomial variants (mhmc/dmhmc) sample any trajectory point proportional to exp(−energy) — is_accepted is always True, and info.acceptance_rate becomes an average trajectory-weight diagnostic, not a reject probability.

We benchmark all four on the Gaussian Process Regression from Section 2 (30 latent variables, 2-D hyperparameter space φ = (log ℓ, log σ_f)). The state_gp and params_gp already produced by Section 2’s window_adaptation are reused directly — no extra warmup needed.

  step_size        = 0.3276
  inv_mass_matrix  = [0.12375075 0.12590666]
LaplaceHMCState fields:         ['position', 'logdensity', 'logdensity_grad', 'theta_star']
LaplaceDynamicHMCState fields:  ['position', 'logdensity', 'logdensity_grad', 'theta_star', 'random_generator_arg']
laplace_hmc           ESS=  341  time=79.1s  ESS/min=  259  M-H acc=0.74
laplace_mhmc          ESS=  316  time=41.2s  ESS/min=  461  traj weight=0.83
laplace_dhmc          ESS=   45  time=75.3s  ESS/min=   36  M-H acc=0.74
laplace_dmhmc         ESS=  183  time=79.9s  ESS/min=  137  traj weight=0.83
Source
<Figure size 1560x480 with 3 Axes>

Reading the results. The acceptance rate panel reveals something important: laplace_hmc can show acceptable M-H acceptance rate yet produce a very low ESS. This could be the classic periodic-orbit pathology — a fixed trajectory length that happens to nearly return to the starting point, causing the chain to take tiny effective steps despite technically accepting. laplace_dhmc breaks the orbit by randomising the step count each transition, so even the same step size yields much better mixing. Multinomial variants (laplace_mhmc, laplace_dmhmc) also avoid the trap: by sampling a random point along the trajectory they implicitly vary the effective displacement, and is_accepted is always True so they never waste a trajectory.

This illustrates the core trade-off: if you can tune num_integration_steps carefully, laplace_hmc is the cheapest option; if you are unsure of the right trajectory length (common in practice), the dynamic and multinomial variants are more robust out of the box.

Quick-reference: when to use each variant.

VariantUse when
laplace_hmcDefault when trajectory length is well-tuned; window_adaptation works directly
laplace_mhmcDrop-in upgrade when M-H acceptance is low or ESS per gradient is poor
laplace_dhmcTrajectory length is hard to tune; randomised steps break periodic-orbit traps
laplace_dmhmcCombines both benefits — best robustness for unknown geometry