The past few articles we’ve written covered the building blocks of Bayesian Inference and Statistical Machine Learning from the ground up: Starting with Bayes’ Theorem, before learning how to set up models using PyMC, then reasoning about their causal structures and the data generating processes, and using them to inform our modelling process for real data. With each Bayesian model we built, we’ve had to formally notate our prior knowledge about our parameters’ probability distributions before fitting them to our data. How the model generates a posterior distribution from our priors and likelihoods/data used to feel completely abstract. But now ‘used to’ here is the key phrase that’s a thing of the past. It’s finally clicked.
With our previous article exploring the earliest variants of the Markov Chain Monte Carlo algorithm, we’ve been able to peel back some of that abstraction. However, the mechanisms that modern probabilistic frameworks like
The Hamiltonian Monte Carlo (HMC), for some context, is a variant of the Markov Chain Monte Carlo (MCMC)** **family of sampling algorithms that’s quietly running as the computational engine behind a significant share of today’s Bayesian models. Although it’s often assumed that MCMC is exclusive to Bayesian Statistics, its roots actually originate from the field of physics where it was originally developed to simulate the behaviour of the hydrogen bomb. MCMC algorithms, particularly HMC, have since had a positive impact on various diverse fields outside of Bayesian Statistics such as in epidemiology where it played a key role in modelling the effects of interventions in Europe during the spread of COVID-19 (Flaxman et al., 2020). In fact, we’ve already utilized HMC (without knowing it) in all our previous Bayesian Regression articles whenever we generated posterior distributions on total hectares burnt in a wildfire.
In this article, we’ll take a close look as to how HMC conceptually generates a posterior distribution, cover the requirements necessary to get the algorithm running, and then apply these concepts by sampling models of varying complexity. Let’s get started!
Before we hit the ground running on Hamiltonian Monte Carlo, let’s take a step back and revisit the first iteration and the conceptual foundation of Markov Chain Monte Carlo — the Metropolis-Hastings algorithm. To quickly recap, the goal of the Metropolis-Hastings algorithm is to obtain samples from multi-dimensional probability distributions where sampling them directly is too difficult and complex. You might then ask: How then do MCMC-class algorithms avoid having to compute a highly complex, multi-dimensional distribution directly? Metropolis-Hastings, and other variants, accomplish this by constructing a random walk through parameter space where, over many draws, the visiting frequencies match the shape of the distribution they’re trying to sample from. Reviewing the basics of Metropolis-Hastings gives us the architectural foundation that HMC inherited and highlights how it resolved its past limitations. If you want a more in-depth refresher on this topic, feel free to navigate here.
Here’s the pseudocode for how the Metropolis-Hastings MCMC algorithm works:
1.Proposal Step: From your current position, propose a direction to move to and sample. The proposal needs to be generated randomly, such as by rolling a die.
2. Acceptance Ratio ( p move
- Note that the earliest iteration of this step in 1953 assumed asymmetric proposalwhere the probability of proposing a move in one direction equaled the probability of proposing the reverse (i.e. a fair die). However, Professor W.K. Hastings (1970) improved this step by adding acorrection factor,to account forq,asymmetric proposalswhich are instances where some proposals were structurally more likely than others (i.e. a loaded die). This matters in practice whenever a parameter has a natural constraint, such as with a standard deviation where its parameter values must be strictly positive.
3. Accept/Reject Step: If the acceptance ratio is greater than or equal to 1 (** pmove ≥ 1**), that means the proposed position is more probable than the current one so we can move there with certainty. And if
- Ifp, then we can move to the proposed position. But ifmove>up, then we must stay put. The lower thatmove≤upis, the more likely the draw fails which therefore means we must stay put.move
For all innovations that Metropolis-Hastings brought to the table in terms of sampling from complex probability distributions, it wasn’t able to keep up as data and raw computational power scaled over time. Metropolis-Hastings, along with older variants of MCMC like Gibbs Sampling, hit a scaling wall when sampling from models with hundreds, let alone, thousands of parameters because it could not navigate a model’s sampling space efficiently. This was in part due to the fact that these algorithms had no sense of a posterior’s overall shape. Older variants would essentially explore a sample space blindfolded because their proposal mechanism in step 1 was random without any sense of where the high-probability regions were.
On the other hand, Hamiltonian Monte Carlo was able to overcome this limitation and sampled high-dimensional distributions more efficiently by taking advantage of their gradient information to guide its proposal mechanism. To understand how it works, we need to think of the algorithm as a simulation of a “Hamiltonian particle” that behaves similarly to a marble moving through an abstract landscape with varying peaks and valleys. However unlike a marble which will roll continuously before losing energy and eventually stopping at the lowest point due to friction, our Hamiltonian particle doesn’t have any friction to deal with. In this simulation of Hamiltonian mechanics, the particle’s trajectory also behaves like a ball being tossed in the air. During its flight upwards, gravity will slow its momentum before hitting a momentary “” where the ball/particle hits zero velocity at the peak of its arc before coming back down to Earth. This is exactly what happens to a Hamiltonian particle moving into a low-probability region of the landscape characterized by upward-facing slopes. The particle slows down the way a ball does climbing toward the top of its arc, and speeds back up moving into a high-probability region, the way a ball accelerates falling back toward the ground. The downward-sloping valleys in the landscape are the high probability regions.
However unlike the ball toss, a Hamiltonian particle doesn’t wait for its momentum to permanently settle to zero before we take a sample. Instead, we let the particle travel for a fixed stretch of simulated time, however far the landscape’s shape allows it to travel in that timespan, and wherever it ends up becomes the “sample” we collect for our distribution. A steep uphill section might mean the particle barely gets anywhere while a smooth downhill stretch might carry it a longer way. Either way, the sampling point isn’t chosen by the particle running out of steam, but by us.
Now that we’ve covered our simulation of the Hamiltonian particle, let’s zoom out to see how it relates to the generation of a posterior distribution from a Bayesian model and discover how each individual detail we covered plays a role in this process. To build our intuition of how Hamiltonian Monte Carlo samples a Bayesian model that’s been fitted to some data, let’s first generate a dataset with two columns of 50 random values:
import arviz as azimport matplotlib.pyplot as pltimport numpy as npimport pandas as pdimport pymc as pmimport requestsimport sqlite3import warningsfrom scipy import statsfrom sklearn.preprocessing import StandardScalerwarnings.simplefilter(action="ignore", category=FutureWarning)pd.set_option("mode.chained_assignment", None)%config Inline.figure_format = 'retina'az.style.use("arviz-darkgrid")az.rcParams["stats.hdi_prob"] = 0.89 # sets default credible interval used by arviz
print("Generating our test data... \n\n")np.random.seed(42)# test datareal = stats.multivariate_normal([0, 0], np.identity(2))x, y = real.rvs(50).Tprint(f"x (len = {len(x)}): {x}, \n\n \y (len = {len(y)}): {y}")
Going back to our simulation, the geometric landscape that influences where the Hamiltonian particle is more likely to travel is calculated based on the model we’re sampling from. Notice how Figure 2 illustrates the movement of the particle in a two-dimensional landscape? To put it simply, each parameter in a model actually represents a dimension of that landscape. Although it’s not possible for us to visualize how this process works with a hundred parameter model, we can be confident that the principles for the Hamiltonian particle’s motion remain the same regardless of its landscape’s dimension size. To compute the landscape for our Hamiltonian particle to traverse through, we need to write what’s called a negative log-posterior function**. **The purpose of the **negative log-posterior function is to compute the negative log-probability of the data and its parameters so that the algorithm knows what the “elevation” is for a given set of values at that position. For some that might have been a mouthful so let’s take a step back. Think of every possible combination of parameter values as a point in a landscape, where the number of dimensions equals the number of parameters in your model. With our synthetic dataset, our model has exactly two parameters, ** μx and
Now if we take our generic model from earlier, the negative log-probability function will look like the following where essentially each line in the model is added together. Notice how the first two terms on the left are summed across every data point (as per our rule above) since these represent the log-likelihood of our data, given the current parameters. On the other hand, the last two appear only once each since there’s just a single ** μx **and
And here’s the code form of that equation:
Great, so hopefully we understand how the “landscape” for our Hamiltonian particle is generated based on a simple model. Now, what does this calc_U() function look like if we’re having to sample from a significantly more complicated one? Let’s test that by bringing back a revamped version of our old Multiple Linear Regression that we used to model wildfire size but instead, we’ll upgrade it to a Hierarchical Multi-Level Regression model!
########################## DATA EXTRACTION ##########################db_link = "https://raw.githubusercontent.com/vanislekahuna/wps-labs/main/data/historical_bc_wildfires/bc_wildfires.db"response = requests.get(db_link)with open("bc_wildfires.db", "wb") as f: f.write(response.content)conn = sqlite3.connect('bc_wildfires.db')cursor = conn.cursor()join_query = """WITH weather AS ( SELECT wind_speed_ms, wind_direction_deg, wind_direction, temperature_c AS temp_c_ign_date, humidity_dewpoint_temperature_2m AS humidity_dewpoint_temp_k_ign_date, soil_temperature_level_1 AS soil_temp_lvl1_ign_date, fire_label AS weather_fire_label FROM weather_data),terraclimate AS ( SELECT monthly_mean_temp_c, monthly_total_precip_mm, monthly_mean_humidity_vpd_kPa, monthly_mean_soil_moisture_mm, fire_label AS terraclimate_fire_label FROM terraclimate_bc_weather)SELECT *FROM historical_bc_wildfiresLEFT JOIN weather ON historical_bc_wildfires.FIRELABEL = weather.weather_fire_labelLEFT JOIN terraclimate ON historical_bc_wildfires.FIRELABEL = terraclimate.terraclimate_fire_labelWHERE historical_bc_wildfires.SIZE_HA IS NOT NULLAND historical_bc_wildfires.SIZE_HA > 0AND historical_bc_wildfires.FIRELABEL NOT IN ('1951-R00037', '1951-R00050', '1951-R00060', '1951-R00067', '1951-R00069', '1951-R00070', '1956-R00107', '1958-V00283', '1956-R00160', '1968-R00088', '1985-V70083', '1985-V70088', '1987-V90016', '1990-V50012');"""df = pd.read_sql_query(join_query, conn)wildfire_df = df[ [ "FRCNTR", "FIRELABEL", "IGN_DATE", "FIRE_CAUSE", "SIZE_HA", "wind_speed_ms", "temp_c_ign_date", "humidity_dewpoint_temp_k_ign_date", "soil_temp_lvl1_ign_date", "monthly_mean_temp_c", "monthly_total_precip_mm", "monthly_mean_humidity_vpd_kPa", "monthly_mean_soil_moisture_mm" ] ].copy()print(f"wildfire_df shape: {wildfire_df.shape} \n")
############################## DATA TRANSFORMATION ##############################centres = wildfire_df["FRCNTR"].unique()wildfire_df["size_log1ptransformed"] = np.log1p(wildfire_df["SIZE_HA"])wildfire_df["wind_speed_log1ptransformed"] = np.log1p(wildfire_df["wind_speed_ms"])wildfire_df["soil_temp_c_lvl1_ign_date"] = wildfire_df["soil_temp_lvl1_ign_date"] - 273.15one_hot_encodding = pd.get_dummies(wildfire_df["FIRE_CAUSE"], dtype=int)wildfire_df = pd.concat([wildfire_df, one_hot_encodding], axis=1)large_fire_df = wildfire_df[wildfire_df['SIZE_HA'] >= 100].dropna()print(f"\n large_fire_df shape: {large_fire_df.shape} \n")
###################################### DATA /PRE-PROCESSING ######################################continuous_cols = ["wind_speed_log1ptransformed", "temp_c_ign_date", "monthly_total_precip_mm", "monthly_mean_humidity_vpd_kPa"]X_cont_std = StandardScaler().fit_transform(large_fire_df[continuous_cols])X_mat = np.column_stack([ X_cont_std[:, 0], large_fire_df["Lightning"].values, large_fire_df["Person"].values, X_cont_std[:, 1], X_cont_std[:, 2], X_cont_std[:, 3],])zone_labels = sorted(large_fire_df["FRCNTR"].astype(str).unique())zone_to_idx = {z: i for i, z in enumerate(zone_labels)}zone_idx = large_fire_df["FRCNTR"].astype(str).map(zone_to_idx).valuesprint(f"Value count of generated zone labels: \n{pd.Series(zone_idx).value_counts().sort_index()}")y = large_fire_df["size_log1ptransformed"].valuespy
################################################################# HIERARCHICAL MULTI-LEVEL BAYESIAN REGRESSION MODELLING #################################################################coords = { "zone": zone_labels, "predictor": ["wind_speed", "lightning", "person", "temp", "precip", "vpd"], "obs_id": np.arange(len(y)),}with pm.Model(coords=coords) as hier_model: zone_idx_ = pm.Data("zone_idx", zone_idx, dims="obs_id") X = pm.Data("X_mat", X_mat, dims=("obs_id", "predictor")) # Population-level (grand mean across zones) mu_alpha = pm.Normal("mu_alpha", mu=4.0, sigma=0.5) # log(100+1) ≈ 4.6, our size floor mu_beta = pm.Normal("mu_beta", mu=0.0, sigma=1.0, dims="predictor") # Between-zone SD — magnitude of cross-zone heterogeneity per effect sigma_alpha = pm.HalfNormal("sigma_alpha", sigma=0.5) sigma_beta = pm.HalfNormal("sigma_beta", sigma=0.5, dims="predictor") # Non-centered parameterization z_alpha = pm.Normal("z_alpha", 0.0, 1.0, dims="zone") z_beta = pm.Normal("z_beta", 0.0, 1.0, dims=("zone", "predictor")) alpha_zone = pm.Deterministic("alpha_zone", mu_alpha + sigma_alpha * z_alpha, dims="zone") beta_zone = pm.Deterministic("beta_zone", mu_beta + sigma_beta * z_beta, dims=("zone", "predictor")) mu = alpha_zone[zone_idx_] + (beta_zone[zone_idx_] * X).sum(axis=-1) sigma_y = pm.HalfNormal("sigma_y", sigma=0.5) y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma_y, observed=y, dims="obs_id") idata = pm.sample( draws=2000, tune=2000, chains=4, target_accept=0.95, # higher than default — hierarchical funnels need this random_seed=42, idata_kwargs={"log_likelihood": True}, # needed for LOO below )n_div = idata.sample_stats["diverging"].sum().item()print(f"Divergences: {n_div}")az.summary(idata, var_names=["mu_alpha", "mu_beta", "sigma_alpha", "sigma_beta", "sigma_y"])
With the following equation + code, the important piece here is not to get stuck on the technical complexity of the model itself but instead on the position of each parameter, whether that’s the prior or the likelihood. The take-home point here is that regardless of our model’s complexity, the elevation at any given point in our Hamiltonian particle’s landscape is essentially just a sum of log-probability contributions from every one of its parameters:
def unpack_params(q, J, P): """ Slice a flat parameter vector q into its named pieces. This mirrors what PyTensor does internally inside PyMC -- every model, no matter how many groups or predictors it has, ultimately reduces to one long vector that the sampler walks through. """ idx = 0 mu_alpha = q[idx]; idx += 1 mu_beta = q[idx:idx+P]; idx += P sigma_alpha = q[idx]; idx += 1 sigma_beta = q[idx:idx+P]; idx += P z_alpha = q[idx:idx+J]; idx += J z_beta = q[idx:idx+J*P].reshape(J, P); idx += J*P sigma_y = q[idx]; idx += 1 return mu_alpha, mu_beta, sigma_alpha, sigma_beta, z_alpha, z_beta, sigma_ydef hierarchical_calc_U(q, X, y, zone_idx, J, P, mu_alpha_prior=(4.0, 0.5), mu_beta_prior=(0.0, 1.0), sigma_alpha_prior=0.5, sigma_beta_prior=0.5, sigma_y_prior=0.5): """Negative log-posterior (potential energy) at a single point q.""" mu_alpha, mu_beta, sigma_alpha, sigma_beta, z_alpha, z_beta, sigma_y = \ unpack_params(q, J, P) # HalfNormal parameters have zero density outside their support. # Real NUTS avoids this entirely by sampling on a log-transformed, constrained scale internally. # We stay on the natural scale here. if sigma_alpha <= 0 or np.any(sigma_beta <= 0) or sigma_y <= 0: return np.inf # Non-centered transform (matches the pm.Deterministic lines) alpha_zone = mu_alpha + sigma_alpha * z_alpha beta_zone = mu_beta + sigma_beta * z_beta # Linear predictor mu = alpha_zone[zone_idx] + np.sum(beta_zone[zone_idx] * X, axis=-1) log_lik = np.sum(stats.norm.logpdf(y, loc=mu, scale=sigma_y)) log_prior = ( stats.norm.logpdf(mu_alpha, *mu_alpha_prior) + np.sum(stats.norm.logpdf(mu_beta, *mu_beta_prior)) # one sum per predictor - 6 predictors (6 parameters) + stats.halfnorm.logpdf(sigma_alpha, scale=sigma_alpha_prior) + np.sum(stats.halfnorm.logpdf(sigma_beta, scale=sigma_beta_prior)) # one sum per predictor - 6 predictors (6 parameters) + np.sum(stats.norm.logpdf(z_alpha, 0.0, 1.0)) # one sum per zone - 6 zones (6 parameters) + np.sum(stats.norm.logpdf(z_beta, 0.0, 1.0)) # one sum per zone per predictor - 6 zones x 6 predictors (36 parameters) + stats.halfnorm.logpdf(sigma_y, scale=sigma_y_prior) ) return -(log_lik + log_prior)
The second function that the Hamiltonian Monte Carlo algorithm needs is the gradient of the negative log-probability which is a separate slope measurement for each individual parameter in the model. Going back to our Hamiltonian mechanics simulation, the element there that represents the gradient of the negative log probability is the force acting on the particle at its current location. If we’re reusing the ball toss analogy, the gradient here is functioning as the gravity in the toss. For contrast, while the negative log-probability function is calculating the shape of the landscape and determining where the peaks and valleys are, the gradients add value by acting as the “gravity” force in the simulation. However, unlike real-life gravity where its strength remains constant, the gravity of the gradient in the Hamiltonian mechanics simulation varies by location. The pull is strongest in the steepest terrains and essentially non-existent on flat surfaces. Now coming back to our simple model, since we only have two parameters in the model, ** μx** and
And the associated Python function for this specific, two-parameter model gradient would be:
Ok, now that we’ve covered the most complex parts of HMC with the functions, let’s quickly cover the two settings we need to tune in order to get the algorithm started. The first is the leapfrog steps ( L) which is simply the count of how many steps our Hamiltonian particle takes to complete one trajectory. If you remember earlier in the ball toss analogy, we mentioned how we specify the span of time that the ball spends in the air? Together, these two settings determine how long our particle’s trajectory runs before we stop and record a sample.
The second setting we need to tune is the step size (ϵ) which is the distance the simulation covers in a single leapfrog set. For example with our simple model notation from earlier, we specified that: “The number of leapfrog steps to ** L = 11** and the step size to
An issue that may arise from these two settings is that if we specify a poor combination of the leapfrog steps and the step size then we can run into an issue known as the U-turn problem. The U-turn problem arises when a Hamiltonian particle’s simulated trajectory travels far enough to loop back to its own starting point or the general area it started from. As a consequence of this issue, we’ll end up with consecutive samples that are highly autocorrelated and will require far more computation for the particle to explore its landscape.
Luckily, Hoffman and Gelman (2014) solved this problem over a decade ago now with their proposal of the No-U-Turn Sampler (NUTS). In essence, what the No-U-Turn Sampler (NUTS) does is automatically detect when a trajectory starts to double back on itself and stops the simulation right before that happens, instead of relying on a fixed, manually-tuned number of leapfrog steps (** L**) or step size (
The last requirement for HMC is the starting point which we’ll arbitrarily assign to (** q **= [−0.1, 0.2]). The starting point will kick off our “Hamiltonian particle” simulation which allows us to sample from a posterior distribution based on which areas the particle visits most frequently. Once we plug in our two functions and two settings, with the flexibility to plug in whichever model’s functions we’re sampling from, then our entire Hamiltonian Monte Carlo sampler will look like the following:
def Hamiltonian_Monte_Carlo(U, grad_U, epsilon, L, current_q, *args): q = current_q.copy() p = np.random.normal(loc=0, scale=1, size=len(q)) # random flick - p is momentum current_p = p.copy() # Make a half step for momentum at the beginning p -= epsilon * grad_U(q, *args) / 2 # initialize bookkeeping - saves trajectory qtraj = np.full((L + 1, len(q)), np.nan) ptraj = qtraj.copy() qtraj[0, :] = current_q ptraj[0, :] = p # Code 9.9 starts here # Alternate full steps for position and momentum for i in range(L): q += epsilon * p # Full step for the position qtraj[i + 1, :] = q # Make a full step for the momentum, except at the end of trajectory if i != L - 1: p -= epsilon * grad_U(q, *args) ptraj[i + 1, :] = p # Make a half step for momentum at the end p -= epsilon * grad_U(q, *args) / 2 ptraj[L, :] = p # Negate momentum at end of trajectory to make the proposal symmetric p *= -1 # Evaluate potential and kinetic energies at start and end of trajectory current_U = U(current_q, *args) current_K = np.sum(current_p**2) / 2 proposed_U = U(q, *args) proposed_K = np.sum(p**2) / 2 # Accept or reject the state at end of trajectory, returning either # the position at the end of the trajectory or the initial position accept = False if np.random.uniform() < np.exp(current_U - proposed_U + current_K - proposed_K): new_q = q # accept accept = True else: new_q = current_q # reject return dict(q=new_q, traj=qtraj, ptraj=ptraj, accept=accept)
And often one of the results we like to see immediately after running the sampler is a trace plot to evaluate its preformance. With a healthy trace plot, we generally want to see two things. The first thing to look for within each individual chain is that the particle should fluctuate rapidly and densely around a stable value, not drift steadily upward or downward over time. The second thing to look for across all four chains is that all four chains should converge on that same stable value. Chains that settle into visibly different regions signal a problem, even if each one looks stable on its own. The trace plot below is an example of a high preforming chain that was generated for the complex Hierarchical Bayesian Regression from earlier.
az.plot_trace(idata, var_names=["mu_alpha", "mu_beta", "sigma_beta"], compact=True)plt.gcf().suptitle("Convergence Diagnostic Trace Plot")plt.tight_layout()
We set out to reverse-engineer the machine behind modern Bayesian inference and by now, we’ve succeeded in taking it apart! Let’s recap what we accomplished. What began as an abstract black box that somehow turned priors and likelihoods into a posterior distribution is a black box no longer. We then traced the lineage from Metropolis-Hastings’ blind, random walk through parameter space, to the moment gradient information transformed that walk into something far more deliberate: A Hamiltonian particle, gliding through a landscape shaped by our own model’s mathematics, accelerating into high-probability valleys and decelerating up low-probability slopes. Lastly, we built the two functions that define that landscape and the forces acting on it, tuned the two settings that govern how far our particle travels, and watched all five pieces come together into a single, working sampler.
What makes this worth understanding isn’t just the mechanics themselves: It’s the foundation with which we’re using Bayesian Inference to drive real-world impact. Every time pymc.sample() runs, this sophisticated, and highly robust, simulation of a “Hamiltonian particle” is what's quietly happening underneath it all. Hamiltonian Monte Carlo isn't a niche statistical curiosity but rather, the quiet computational engine behind a significant share of the inference done in Bayesian statistics today. For example, we noted in the intro how it played a direct role in modeling the effects of interventions across Europe during COVID-19 which helped shape real public health decisions during a genuine crisis (Flaxman et al., 2020).
We even saw its reach firsthand, too. The same skeleton that generated a posterior distribution for a toy two-parameter model was able to scale, without needing to be reinvented, to a complex Hierarchical Bayesian Regression model spanning dozens of parameters across BC’s Fire Centre zones. With our HMC quietly working in the background, our Hierarchical Bayesian Regression model was able to estimate how wind speed, lightning, precipitation, and a handful of other atmospheric variables could shape the size of a BC wildfire, with the uncertainty in that estimate quantified honestly rather than papered over. Our toy example of HMC scaled to a complex machine running on real data, producing a real answer to a real question!
We came into this article treating the posterior as something that simply appeared. Now we can now proudly say we’re leaving it with a working understanding of exactly how it gets there though one leapfrog step, one gradient, and one accepted proposal at a time.
Thank you for reading through to the very end. This was an amazing journey I’m glad we were able to travel through together.
Reverse-Engineering Hamiltonian Monte Carlo: The MCMC Engine Behind Modern Bayesian Inference was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.