1. Setup: HMC for the 8-schools hierarchical model

The model is the classic educational testing / 8-schools example:

  • There are 8 schools, indexed by j = 1,…,8.
  • Each school has an observed treatment effect estimate $y_j$ with known standard error $\sigma_j$.
  • The model (in Chapter 5) is:
    • Likelihood: $y_j \mid \alpha_j \sim N(\alpha_j, \sigma_j^2)$
    • Hierarchical prior: $\alpha_j \mid \mu, \tau \sim N(\mu, \tau^2), \quad j=1,\dots,8$
    • Hyperpriors on $\mu$ and $\tau$ (not detailed here, but they’re such that the posterior is proper).

To avoid conflict with earlier notation, they now call the school-level effects αⱼ.

The full parameter vector is:

$\theta = (\alpha_1, \dots, \alpha_8, \mu, \tau)$

so the dimension is:

$d = 8 + 1 + 1 = 10$

HMC is not required for this small model (a Gibbs sampler with parameter expansion works very well), but this example is used to demonstrate how to set up and tune HMC step by step on a simple hierarchical model.

2. Gradients of the log posterior

For HMC, we need the gradient of the log posterior with respect to all parameters.

We write:

$\log p(\theta \mid y)$

and we need:

  • $\dfrac{\partial \log p(\theta \mid y)}{\partial \alpha_j}$​, j = 1,…,8
  • $\dfrac{\partial \log p(\theta \mid y)}{\partial \mu}$
  • $\dfrac{\partial \log p(\theta \mid y)}{\partial \tau}$

Using the normal likelihood and normal hierarchical prior, they derive:

  1. Gradient with respect to each αⱼ:
    • $\frac{\partial \log p(\theta \mid y)}{\partial \alpha_j} = – \frac{\alpha_j – y_j}{\sigma_j^2} – \frac{\alpha_j – \mu}{\tau^2}, \quad j = 1,\dots,8$
      • The first term comes from the likelihood:
        • $y_j \mid \alpha_j \sim N(\alpha_j, \sigma_j^2)$
      • The second term comes from the prior:
        • $\alpha_j \mid \mu, \tau \sim N(\mu, \tau^2)$
  2. Gradient with respect to µ:
    • $\frac{\partial \log p(\theta \mid y)}{\partial \mu} = – \sum_{j=1}^{J} \frac{\mu – \alpha_j}{\tau^2}$ where J = 8 here.
    • This comes from the fact that $\mu$ appears only in the prior $N(\mu, \tau^2)$ for the αⱼ.
  3. Gradient with respect to τ:
    • $\frac{\partial \log p(\theta \mid y)}{\partial \tau} = – \frac{J}{\tau} + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^3}$
      • The $-J/\tau$ term arises from the normalizing constant of the normal prior on αⱼ.
      • The sum term comes from the quadratic part $(\alpha_j – \mu)^2/(2\tau^2)$.

Debugging the gradients

As a sanity check, they also compute the gradients numerically using finite differences:

  • For each component of θ:
    • Perturb that parameter by ±0.0001,
    • Approximate the derivative by:
      • $\frac{\log p(\theta + \delta e_k) – \log p(\theta – \delta e_k)}{2\delta}$ where eₖ is the unit vector in coordinate k.

They verify that:

  • The analytic gradient and the finite-difference gradient agree to several decimal places.
  • Once verified, they use the analytic one in HMC, because it is much faster.

3. Choosing the mass matrix $M$ for the momentum φ

Recall:

  • HMC introduces momentum variables $\phi \sim N(0, M)$, with the same dimension as θ.
  • M is called the mass matrix.

Goal:

  • Scale M so that it roughly matches the typical posterior scales in each dimension.
  • This helps trajectories move efficiently and not be dominated by huge scale differences.

In this example:

  • They look at the data and conclude a typical scale of about 15 for all parameters.
  • So they set: $M = \text{Diag}(15^2, \dots, 15^2)$, i.e. a diagonal mass matrix where every diagonal entry is $15^2$.

This is crude but sufficient to avoid gross mismatch in scaling across dimensions.

4. Starting values (initialization)

They run 4 parallel chains.

For starting values:

  • For each chain, draw each of the 10 parameters independently from: $N(0, 15^2)$

This is consistent with the rough scale they assigned to the mass matrix and ensures the initial states are dispersed over a reasonable region of parameter space.

5. Tuning ε (step size) and L (number of leapfrog steps)

Instead of fixing ε and L, they randomize them around central values:

  • Choose central values $\epsilon_0$​ and $L_0$​.
  • At each iteration:
    • Draw ε uniformly from $(0, 2\epsilon_0)$.
    • Draw L from a discrete uniform distribution on $\{1, \dots, 2L_0\}$.

This “jittering”:

  • Is not theoretically optimal but is simple and safe (it doesn’t depend on past outcomes, so detailed balance is preserved).
  • Helps the algorithm take variable-length trajectories and step sizes.

Initial setup:

  • They follow the general guideline that $\epsilon_0 L_0 = 1$.
  • They choose:
    • $\epsilon_0 = 0.1$,
    • $L_0 = 10$,
      so $1 \times 10 = 1$

First, they run 4 chains for 20 iterations just to confirm the code runs and doesn’t crash.

5.1 First tuning run (100 iterations)

They run 4 chains for 100 iterations and check:

  • Posterior inferences look plausible (no absurd outliers).
  • But the chains are not close to convergence:
    • Several Gelman–Rubin $\hat{R}$ values greater than 2.

Average acceptance rates per chain are:

  • 0.23, 0.59, 0.02, 0.57

The average rates are far below 65%, especially that chain with 0.02, so they conclude:

The step size ε is too large; the trajectories are too aggressive.

5.2 Second tuning run (smaller ε₀, larger L₀)

They adjust:

  • Decrease $\epsilon_0$​ to 0.05,
  • Increase $L_0$​ to 20,

so that:

  • $\epsilon_0 L_0$​ remains 1 (the same total “path length”).

Run 4 chains for another 100 iterations. Now the acceptance rates are:

  • 0.72, 0.87, 0.33, 0.55

So they are higher, closer to the desired ~65% level, but:

  • Chains still have not mixed well (not enough iterations).

5.3 Longer run (1000 iterations)

They increase the number of iterations to 1000 with these new settings.

Results:

  • Chains are now close to convergence, with: $\hat{R} < 1.2 \text{ for all parameters}$
  • Acceptance rates stabilize at:
    • 0.52, 0.68, 0.75, 0.51

These are near the rule-of-thumb target (~65%).

5.4 Final long run (10,000 iterations)

They then run 4 chains × 10,000 iterations using these tuned settings.

Outcome:

  • Approximate convergence achieved: $\hat{R} < 1.1 \text{ for all parameters}$

Conclusion:

  • For this small model, HMC works fine, but it’s not necessary.
  • A well-parameterized Gibbs sampler is simpler and just as good here.
  • However, for larger and more complex models, Gibbs or random-walk Metropolis can become too slow, while HMC can still move efficiently in high dimensions.

6. Transforming τ to log τ

Because τ is constrained:

  • $\tau > 0$,

HMC on τ directly can:

  • Propose negative τ during leapfrog steps,
  • Hit the boundary and waste iterations.

A common remedy is to transform to an unconstrained parameter:

  • Let: $\tilde{\tau} = \log \tau.$
  • Then $\tilde{\tau} \in \mathbb{R}$, unconstrained.

To work in log τ, we must adjust:

  1. Redefine θ
    • Replace τ with log τ.
    • New parameter vector:
      • $\theta = (\alpha_1, \dots, \alpha_8, \mu, \log \tau)$
  2. Adjust the posterior density with the Jacobian
  3. Because:
    • $\tau = e^{\log \tau}$,
    • $d\tau = \tau \, d(\log \tau)d$,
    • the jacobian determinant of the transformation is $\tau$. So:
    • The posterior density in terms of the new θ becomes:
      • $p_{\text{new}}(\theta \mid y) = p_{\text{old}}(\alpha_1,\dots,\alpha_8,\mu,\tau \mid y)\ \times \tau$
    • In log terms:
      • $\log p_{\text{new}}(\theta \mid y) = \log p_{\text{old}}(\alpha_1,\dots,\alpha_8,\mu,\tau \mid y) + \log \tau$
    • So in the code, you add log τ to the log posterior.
  4. Adjust the gradient with respect to log τ
    • Before the transformation, we had:
      • $\frac{\partial \log p}{\partial \tau} = -\frac{J}{\tau} + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^3}$
    • After we add $\log \tau$, the derivative with respect to τ becomes:
      • $\frac{\partial}{\partial \tau} [\log p_{\text{old}} + \log \tau] = \left(-\frac{J}{\tau} + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^3}\right) + \frac{1}{\tau} = -\frac{J-1}{\tau} + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^3}$
    • But in the new parameterization, the gradient is with respect to log τ, not τ. Using the chain rule:
      • $\frac{\partial \log p}{\partial \log \tau} = \frac{\partial \log p}{\partial \tau} \cdot \frac{\partial \tau}{\partial \log \tau} = \left( -\frac{J-1}{\tau} + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^3} \right) \cdot \tau$
    • which simplifies to:
      • $\frac{\partial \log p(\theta \mid y)}{\partial \log \tau} = – (J – 1) + \sum_{j=1}^{J} \frac{(\mu – \alpha_j)^2}{\tau^2}$
    • This matches the formula given.
  5. Change the mass matrix
    • The typical scales are now:
      • For $\alpha_1, \dots, \alpha_8, \mu$: still around 15.
      • For $log \tau$: they use scale 1.
    • So they set:
      • Mass for $\alpha_j$ and $\mu$:15
      • Mass for $\log \tau$: 1
  6. Change initial values
    • For the first 9 parameters ($alpha_1, \dots, \alpha_8, \mu$):
      • $\sim N(0, 15^2)$
    • For $\log \tau$:
      • $\sim N(0,1)$

After these changes, HMC proceeds exactly as before, but in the transformed, unconstrained space. They again start with ε = 0.1, L = 10, and then adjust to reach a reasonable acceptance rate.

7. Big picture

This example shows, in a concrete hierarchical model:

  • How to:
    • Derive analytic gradients,
    • Choose an initial mass matrix,
    • Initialize multiple chains,
    • Tune ε and L based on acceptance rates and $\hat{R}$,
    • Transform constrained parameters (τ > 0) to unconstrained ones (log τ),
    • Adjust log density and gradients for the Jacobian.

In this particular case:

  • HMC is educational, not necessary.
  • For bigger, more complex hierarchical models, HMC can dramatically outperform Gibbs and basic Metropolis in terms of effective sample size per unit computation.