Energy Based Models for AI (Part 3): The Sampler

In this part, we will discuss on the sampling. A crucial component to an energy based learning.

The most important part in an Energy Based Learning is the sampling. Which is one of the way for us to calculate the ‘loss’ of our model and train the model.

Why sampling? Well, the idea behind Energy Based Learning comes from the Gibbs Free Energy Principle. Gibbs Free Energy states that each particle that containes energy will in excited state (high state) until it all of the energy it has is spent. Then it becomes rested in the low state.

So, Hinton and LeCunn are taking that inspiration and applied it to the world of data. The bad data configuration are in the high energy state, and the good data configuration is located in the low energy state, the valley so to speak.

We can use Gibbs Free Energy equation, and write the energy based learning equation as below, using the Boltzmann distribution:

where 𝑃(𝑥) is the probability distribution of x. And 𝐸θ(𝑥) is the scalar energy function. Which in this case is the neural network itself. The one that we want to focus right now is the 𝑍(θ), which is what we called as partition function.

Let’s rewrite that Z,

The Z is actually the overall sum of the energy function across all possible states. That is not possible to be calculated, and what mathematician call as intractable. Therefore, the best way to get this is to sample.

But how did we do that? This bring us to the Sampling.

Just the gradient, please

Let us recap what happen during the training of neural network. Specifically during backpropagation t. During that time, the gradient that we get from the loss function is brought back to the beginning of the network, and is used to update the model weights.

The process of updating that weights is called gradient descent, and we can write that as an equation below:

Where 𝜃 is the model weights.As you can see there, there’s a derivative there in the equation. If we remember our high-school mathematics (or Add-Math, if you’re from Malaysia), that derivative can be called as gradient. Why gradient? Because when we solve a derivative plotted on a cartesian coordinates, we get a slope of the straight tangent line touching our curve, or as we like to called, a gradient.

However, that simple equation does not paint the full picture. In a neural network, we are dealing with thousands, or maybe even millions, of weights. This means our single gradient vector has thousands or millions of dimensions to account for.

At the start of training, this gradient is rarely pointing toward the true, final answer. Because the gradient only calculates the slope immediately under its feet, it will point directly toward the nearest shallow dip in the landscape, a trap we call a local minimum.

We can imagine this high-dimensional mathematical space as a massive, rugged mountain with multiple nooks, crannies, and false valleys. Our ultimate goal during training is to navigate this treacherous landscape and safely descend to the absolute lowest base of the mountain, known as the global minimum.

This methaphor of hills and valley almost sounds like energy-based learning right? Yeah, that’s basically the initial motivation of it. We can apply the same logic to our energy-based learning, especially in building our sampler. To find the lowest energy point, we can use the gradient.

Understanding the Gradient Ascent

So, we now know to use the gradient as sampling criteria. With this information we can build our sampler. Let’s take a look back at our Boltzmann Distribution equation:

where the denominator is our partition function. Using logarithmic rule of mathematics, we can write the equation as below:

The left part of our equation is easy to solve, that’s the energy cost from the real data configuration. What we want is the right term, that is the part where we need to use our sampler.

So, for sampling we need to add gradient term. Let’s do that.

Once we put the gradient there, we can simplified our equation to

Note: There is actually a massive mathematical proof under the hood that gets us to this elegant approximation. We can’t just slap a gradient on the equation and magically hope that terrifying summation term disappears. There is a whole trail of calculus lemmas to trace and understand here. But since I figured most of you are already reliving your university exam trauma at this point, I’m going to skip the gory details for your sanity (haha).

What is important for us to take away is how we use this to update our weights. One crucial twist to note here: in Energy-Based Learning, we aren’t using traditional gradient descent to minimize a loss function. Instead, we are using gradient ascent to maximize our log-likelihood.

Our goal here is a dual-action tug-of-war. We want to use this gradient to actively push up the energy of our fake, negative samples (flattening the false valleys) while simultaneously pulling down the energy of our real data samples (digging out the true valleys).

Then, we we update our weights, the equation would be as below:

where 𝜂 is our learning rate. The equation above tell us the full story of energy based learning. The left term is the correct data, the right term is fantasy state. Our goal is to update our model weights, until the right terms will be equal with the left term. Which means the fantasy state is equal to the real state.

That’s the beauty of it.

Alright, I think I’m gonna stop here. The math is too long already, and I don’t want to burden who’s reading this. We shall continue next time with implementation.