Learning Energy Based Models¶
Energy based models are models that place an unnormalized probability distribution over a space. The density of this probability distribution is proportional to the negative exponent of the "Energy" associated with a state.
That is:
$$ P(x) =\frac{1}{Z}e^{-E(x)} $$Where $Z$ is called the partition function defined as:
$$ Z = \int e^{-E(x)}dx $$Most of the time we're intersted in learning some model via maximum likelihood or some other related method such as map or bayesian sampling. In this case we can paramaterize our model with parameters $\theta$. We can then write the probability distribution as
$$ P(x | \theta) =\frac{1}{Z(\theta)}e^{-E(x,\theta)} dx$$Where the partion function now depends on theta:
$$ Z(\theta) = \int e^{-E(x,\theta)} dx $$So far so good, if our model was discrete and compute wasn't an issue, we could calculate $Z$ and then $P(x | \theta)$ and find the theta that maximises our model. But suppose $E(x,\theta)$ was a Neural Network, or more generally, suppose we can differentiate $E(x,\theta)$ w.r.t $\theta$. We can ask, what is the gradient w.r.t $\theta$ of $P(x | \theta)$ ?
We can compute this:
$$\nabla log(P(x|\theta))$$$$= \nabla \frac{1}{Z(\theta)} e^{-E(x)}$$$$= \int P(x'|\theta)\nabla E(x',\theta) dx' - \nabla E(x,\theta)$$In other words, we can estimate the gradient by computing the expectation of the gradient w.r.t samples taken from the distribution implied by the energy function. The only peice missing from this puzzle is how to do that. Fortunately this is a well studied problem and can be solved by using Markov Chain Monte Carlo.
This is essentially what Open AI did in their paper on EBMs. They used a gradient based sampler called Langevin dynamics to gather samples to learn the model. It's also the operating principle behind Hinton's now ancient RBMs. With block gibbs sampling replacing SGLD.
The following demonstrates the principle with a NN working on MNIST
#importing dependencies
import torch as t
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as f
from torch import FloatTensor as FT
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import time
from torchvision.datasets import MNIST
from tqdm import tqdm
from scipy.stats import linregress
%matplotlib inline
#Loading MNIST dataset
PATH_TO_MNIST = "./"
mnist_ds = MNIST(PATH_TO_MNIST,train=True)
mnist = np.stack([np.array(x[0]) for x in mnist_ds])
mean = np.mean(mnist)
std = np.std(mnist)
scaled_mnist = (mnist-mean)/std
#Utility for displaying mnist image
def disp_img(img):
plt.imshow(img)
plt.show()
disp_img(scaled_mnist[0])
Here we define the HMC sampler as described in Bishop's PR&ML. Note that this is not a maximally efficent implementation as the energy function is evaluated twice at the same location.
Any MCMC sampler should work, a 6th order HMC Sampler is used because the higher order integration makes it possible to use a larger step size.
def grad(e,x,ret_e=False):
with t.enable_grad():
x.requires_grad = True
x.grad = None
energy = e(x)
t.sum(energy).backward()
g = x.grad
x.requires_grad = False
if ret_e == True:
return g,energy
else:
return g
return g
def symplectic_operator2(e,dt,x,z,x_grad=None):
z.requires_grad = False
if x_grad is None:
x_grad = grad(e,x)
z -= (dt/2)*x_grad
x += (dt)*z
x_grad = grad(e,x)
z -= (dt/2)*x_grad
return x,z,x_grad
def symplectic_operator6(e,dt,x,z,x_grad=None):
#https://doi.org/10.1016/0375-9601(90)90092-3
#Table 1: Solution A
w_list = [0.78451361047756e0, 0.23557321335935e0, -0.117767998417887e1]
w0 = 1-2*(sum(w_list))
w_sequence = w_list + [w0] + w_list[::-1]
for w in w_sequence:
x,z,x_grad = symplectic_operator2(e,w*dt,x,z,x_grad)
return x,z,x_grad
def hmc6(x,e,dims,dt=0.1,n_step=10,reject=True):
z = t.normal(mean=t.zeros_like(x))
x_init = x.clone()
if(len(dims) != 0):
h_init = (1/2)*t.sum(z**2,dims) + e(x)
else:
h_init = (1/2)*z**2 + e(x)
x_grad = None
for i in range(n_step):
x,z,x_grad = symplectic_operator6(e,dt,x,z,x_grad)
if(len(dims) != 0):
h_final = (1/2)*t.sum(z**2,dims) + e(x)
else:
h_final = (1/2)*z**2 + e(x)
switch = t.log(t.rand_like(h_init)) < (h_init-h_final)
switch = switch[[slice(None)]+[None for x in dims]]
acc_rate = t.mean(switch.type(t.FloatTensor))
acc_rate = t.mean(switch.type(t.FloatTensor))
if reject:
x_final = t.where(switch,x,x_init).detach()
else:
x_final = x.detach()
return x_final,acc_rate
To test that the sampler is working correctly, we pass it the gaussian energy function, and validate that the samples are drawn from the correct distribution. Note the high acceptance rate even with a large step size.
#Test HMC Sampler
x = t.zeros((100000,))
e = lambda x: (1/2)*x**2
for i in range(100):
x,acc = hmc6(x,e,[],1,10)
if i % 20 == 0:
print(acc)
plt.hist(x,bins=100,density=True,label="Sampler Histogram")
y = np.linspace(-3,3,1000)
plt.plot(y,np.exp(-y**2/2)/np.sqrt(2*np.pi),label= "Gaussian Density")
plt.legend()
plt.show()
Here we define the model, which is just a function from an image to an energy. There are two features of note.
-The model has been defined in such a way that the energy has a minimum at 0. This isn't nessesary, but it seems to stabalize training and prevents the model from blowing up.
-The model imposes a penalty on the L2 norm of the input, this is to ensure that the model is always normalizable and that the partition function is well defined. Again, this is not stricly nessesary since unormalizable models have low loss according to the objective function being optimized, but again, it helps stabilize training.
#Define the model
class EBM(nn.Module):
def __init__(self,h,y):
super(EBM, self).__init__()
self.l1 = nn.Linear(28*28,h*2)
self.l2 = nn.Linear(h*2,h)
self.l3 = nn.Linear(h,y)
def forward(self,x):
x = x.view(-1,28*28)
h1 = f.elu(self.l1(x))
h2 = f.elu(self.l2(h1))
out = self.l3(h2)
a = t.sum(out**2,[1])
b = t.sum((1/2)*x**2,[1])
energy = a + b
return energy
Here we train the model by interleaving model update and model sampling steps. A few things are notable.
-A very high ratio of sampling steps to model update steps, this is because for learning to occur samples should be drawn from the distribution similar to that of the current model. Thus when the model is updated, the samplers need to return to equlibrium.
-With the exeption of the start. The "loss" stays near zero during the entirety of the training procedure, this is a good thing, as it implies that the samples and the data both have similar probablity according to the model. This indicates that the model is well calibrated and that the samples are being drawn from the equilibrium distrubtion of the model.
-The sampler acceptance rate drops steadily over the course of training. This is because the energy landscape becomes more jagged as the model learns to concentrate probability mass. Thus for a constant step size, the sampler will be prone to making more errors.
-If the model changes faster than the sampler can keep up, such that the samples no longer come from the equilibrium distribution then the the loss will begin to drop quickly, locking up the sampler (acceptance rate goes to 0) and then the loss will implode to large and negative values.
As a result, training models in this way requires a very large amount of compute with this model taking a few hours to train to get these decent, but not stellar results.
#Define training procedure
BS = 256
N_CHAINS = 20000
N_ITER = 1000000
DT = 1e-2
N_STEP = 10
GRAD_CLIP = 1
DE = 0
def train_ebm():
chains = t.normal(t.zeros(N_CHAINS,28,28),std=1).cuda()
ebm = EBM(1000,10).cuda()
data = t.FloatTensor(scaled_mnist).cuda()
#Add noise to prevent degenerate energy functions
noise_data = t.normal(data,std=1e-2)
opt = optim.Adam(ebm.parameters(),lr=1e-4)
loss_list = []
acc_list = []
for i in range(N_ITER):
try:
with t.no_grad():
pos_perm = t.randperm(data.size(0))
pos_idx = pos_perm[:BS]
pos_samples = noise_data[pos_idx]
neg_perm = t.randperm(chains.size(0))
neg_idx = neg_perm[:BS]
neg_samples = chains[neg_idx]
neg_samples,acc = hmc6(neg_samples,ebm,[1,2],DT,n_step=N_STEP)
chains[neg_idx] = neg_samples
loss = t.sum(-(ebm(neg_samples)-ebm(pos_samples)))/BS
opt.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(ebm.parameters(), GRAD_CLIP)
opt.step()
loss = loss.detach().cpu().numpy()
acc = acc.detach().cpu().numpy()
#if i % 10000 == 0:
#print(i,loss,acc)
loss_list.append(loss)
acc_list.append(acc)
except KeyboardInterrupt:
break
return loss_list,acc_list,chains
def plot_ebm(data):
loss_list,acc_list,chains = data
plt.title("Loss")
plt.plot(loss_list)
plt.show()
plt.title("Loss FFT")
plt.plot(np.log(np.abs(np.fft.rfft(loss_list))))
plt.show()
plt.title("Acceptance Rate")
plt.plot(acc_list)
plt.show()
for j in range(20):
img = chains[j].cpu().numpy()
disp_img(img)
data = train_ebm()
The main takeaway is that neural EBM training is possible but that a very large ratio of sampling steps to training steps is nessesary. Both this modela and the OpenAI work used a 60:1 ratio. To make this an effective method, either compute needs to be scaled up or a more efficient sampler must be used.
plot_ebm(data)
