Python: Hamiltonian Monte Carlo from scratch | by Jake | Towards Data Science

If your experience with Markov Chain Monte Carlo (MCMC) methods has been anything like mine, it’s been an endless search for answers. Every time you get an answer to one question, a dozen more pop up! You’re continuously learning how much you don’t know. I recently accepted a full time role as data scientist at an organization that makes extensive use of Bayesian Statistics and by extension, MCMC methods. So I decided to go to war with this branch of mathematics once and for all! For a detailed review of Metropolis, Metropolis-Hastings, and Hamiltonian Monte Carlo, please visit my public Google Colab Notebook :)
I’ll revisit the key ideas behind Metropolis-Hastings (MH) before diving into the new material because Hamiltonian Monte Carlo (HMC) is an extension of MH.
Metropolis-Hastings is a glorified random walk. You need four elements: A starting point, a target distribution, a proposal distribution, and an impartial judge (a random event.) The target distribution can be any distribution you’d like to sample from; all you need is the probability density function (PDF). Likewise, the proposal distribution can be any PDF (although the math is simpler if the distribution is symmetric.) The whole idea is start at the random point using the proposal distribution. Append it to an array. Feed this point into the target’s PDF to get the probability density (aka likelihood). Now, generate a disturbance via the proposal distribution, add it the the current location, assigning this variable a new name, and evaluate its likelihood, as well, via the target’s PDF. Compare the likelihoods
acceptance = target_PDF(proposed)/target_PDF(current)
This number will fall anywhere in the range [0, inf). If it’s less than 1, there’s an x% chance that you move. Anything above 1 is guaranteed movement. This is where the “impartial judge” comes in. Observe a random event (typically a number in [0,1] sampled from a uniform distribution.) If this number is less than or equal to acceptance, you move to the proposed point. Else, you stay at the previous location. Append the winner to the array. And that’s it to the Metropolis algorithm.
Things are slightly more complicated if the proposal distribution is asymmetric. If you sample movements from a normal distribution, uniform distribution, etc this step is unnecessary. However, if you use a distribution, such as gamma, poisson, exponential, lognormal, etc. then you’ll need to account for the inherent bias. In the Colab Notebook I linked at the beginning, I use a Beta proposal distribution with a (-0.5 constant term), biased towards 1 (as opposed to symmetric or biased toward 0.) Two things to note: First, the Beta distribution is defined [0,1] meaning we will never observe a negative disturbance, which means that our proposals will increasingly move only one direction. I added the -0.5 bias such that now positive and negative samples can disturbances can be sampled. Second, however, the distribution is still biased towards 1, meaning that it will slowly creep in only one direction. We need to account for this. We do this by getting the likelihoods of the disturbance and the reverse disturbance.
def correct(prop,curr,a=3,b=2):    x0 = curr - prop + 0.5    x1 = prop - curr + 0.5    b0 = beta.pdf(x=x0, a=a, b=b)    b1 = beta.pdf(x=x1, a=a, b=b)    return b0/b1
This correction term (as you’ll see in the Colab notebook) is multiplied by the likelihood ratio; it accounts for the asymmetry of the proposal distribution. (Note, this ratio is called the Hastings-ratio as it’s the only difference between Metropolis and MH algorithms.)
You can sample from most distributions using MH algorithm. However, at the end of the day, it’s fundamentally a random walk. There’s no logic informing how large the jumps should be given the current position. When sampling 1 variable distributions, this isn’t an issue. But as dimensionality increases, the region of high likelihood composes only a fraction of the total area, whereas the moderate and low likelihood regions compose (exponentially) more of the total area. This effect imposes only a mild inefficiency on MH when sampling low dimensional distributions. But as dimensionality increases, MH runs the risk of returning samples that aren’t representative of the target distribution.
The reason is that small jumps are appropriate near the peak(s), because over-exploration of this region, relative to its likelihood, is low risk; however, as we move towards the tails, over-exploration of these areas, relative to their likelihoods, becomes a sizable risk. (Many samples from one tail and none in the other isn’t desirable, for context.) So...how do we propose jumps with likelihood in mind?
Physics has the answer — hurray! For the remainder of this post, we’re not going to view locations as a random walk. Rather, we’re going to view the absolute peak (or peaks) of a distribution as a planet. And a satellite will orbit this planet, collecting samples. Hamiltonian mechanics uses differential equations relating kinetic energy to potential energy. These differential equations have really tricky exact solutions but friendly approximate solutions using leapfrog integration. Integration of our differential equations will give us the path which our satellite will follow. These differential equations (or Hamiltonian equations) define the energy of a system in terms of kinetic and potential energy. When you throw a ball up into the air, its kinetic energy is replaced by potential energy. As the ball falls, its potential energy is replaced by kinetic energy. Hamiltonian equations define the relationship between position and momentum:
T = timeQ = position P = momentum K = kinetic energyV = potential energydQ/dT = PdP/dT = dV/dQ 
Note the above equations are derived in the context of statistics. The physics world is necessarily, far more complicated.
Of note, our differential equations are dQ/dT, the change in position wrt (with respect to) time and dP/dT, the change in momentum wrt time, which evaluates to the change in potential energy wrt position. We can define position in terms of the PDF, itself. And so, changes in momentum are a function of position in the PDF, namely distance from the peak. Note that the normal distribution’s gradient is steepest at inflection points between the peak and the tail(s), respectively. Conversely, the negative log normal distribution’s gradient is steepest at the tails and approaches (and reaches) 0 as it approaches the peak. As we discussed earlier: We want our jump size ro be proportional to the location in the distribution. Near the peak, we want to take small jumps. As we approach the tails, we want to take progressively bigger and bigger jumps. It turns out that the negative-log normal distribution gradient accomplishes this effect perfectly. (Whereas the normal distribution gradient would propose small jumps from both the peak and the tails, which we don’t want!)
The algorithm breaks down into four parts:
Set up: Take the previous position and copy, such that you have q0 and q1. Randomly sample a momentum from N(0,1) and copy, such that you have p0 and p1. Find the gradient of the PDF with respect to position -(x-mu)/sigma^2 for a single variable gaussian.Leapfrog: Use leapfrog integration to update q1 and p1 (ie Hamiltonian motion or a particle.) In practice, this is the most sensitive part; small adjustments can induce unstable behavior.MH Dance: Lastly, multiply p1 by (-1), this ensures reversibility (ie q1 can be reached from q0 given momentum p0 AND q0 can be reached from q1 given momentum -p1). This gives us the information we need for the metropolis-hastings “dance”. Keep in mind we’re using negative log probabilities, so the math is all in terms of addition and subtraction. Then accept/reject movement.Repeat for fixed number of iterations.
Set up: Take the previous position and copy, such that you have q0 and q1. Randomly sample a momentum from N(0,1) and copy, such that you have p0 and p1. Find the gradient of the PDF with respect to position -(x-mu)/sigma^2 for a single variable gaussian.
Leapfrog: Use leapfrog integration to update q1 and p1 (ie Hamiltonian motion or a particle.) In practice, this is the most sensitive part; small adjustments can induce unstable behavior.
MH Dance: Lastly, multiply p1 by (-1), this ensures reversibility (ie q1 can be reached from q0 given momentum p0 AND q0 can be reached from q1 given momentum -p1). This gives us the information we need for the metropolis-hastings “dance”. Keep in mind we’re using negative log probabilities, so the math is all in terms of addition and subtraction. Then accept/reject movement.
Repeat for fixed number of iterations.
You might have noticed in step-1 that we sampled momentum from a gaussian. What’s that all about? Using our satellite metaphor from before, Hamiltonian motion will guide our orbital path around a specific likelihood (visiting positions associated with the same likelihood.) However, we’re not interested in exploring just one likelihood, we’d like to explore them all. In order to accomplish this effect, we sample momentum “kicks,” which can cause our satellite to jump or fall orbital distances (aka visit different likelihoods.)
And let’s take a look at performance!
If you’d like to run my code, please hop onto the Google Colab link from the beginning. You’ll only have viewer access, not editor access — but you can copy to your own Google Drive for experimentation, download the iPython Notebook file, etc. If you play around with the code, you’ll notice that path length and step size are very sensitive hyper-parameters. The slightest adjustment might result in nonsensical samples. This is the curse of gradient-based MCMC samplers. However, there’s hope — newer methods, such as No-U-Turn-Sampler (NUTS) builds upon HMC by dynamically choosing path length and step size.
But why is this a problem? Well, from our hypothetical satellite orbit example, Hamiltonian equations will guide the motion of the satellite, “orbiting” a particular likelihood, meaning it proposes new locations near similar likelihoods. But the orbital path around any given likelihood varies in length. At higher likelihoods, the circumference is quite small but larger when orbiting lower likelihoods (closer to the tails.) We can visually confirm this by simply looking at the gaussian distribution. What happens when me make a 360 orbit around a given likelihood? We propose the same point we’re currently at. This is called a “U-Turn” in HMC lingo, meaning that we come back to where we’ve just been. Dynamic adjustments to path length and step size control the extent of integration (the orbital path.) And as you’ll see if you play with the code a bit, these hyper-parameters are quite sensitive — hence you’ll likely use NUTS over HMC in practice.
I’ve got to give credit where it’s due. Colin Carroll, a software engineer at Google, wrote a more scalable implementation of HMC. (Mine can only be used to sample one-variable gaussians.) The scalability came at a price — interpretability. I found that to really “get” the concept, I had to dumb it down and build one from scratch. Hopefully, you’ve enjoyed the read!
Thank you for reading — If you think my content is alright, please subscribe! :)