Micropolis I

Micropolis II: Minibatch MH

In Part I we established the most basic sanity test for a machine learning algorithm---that it is able to fit data, at all. In this installment, we'll implement a big pratical improvement over that basic algorithm: minibatches.

Recall that our Metropolis-Hastings optimizer works by randomly perturbing the model parameters, computing the loss for the new parameters, and then accepting or rejecting the new parameters with probability $e^{({\tt old\_loss}\, -\, {\tt new\_loss})/T}$ for temperature $T$. So, accepting with probability 1 if the new loss is better and approaching 0 as the new loss gets worse.

Our original optimizer computed the loss using the full dataset, which worked great for our toy dataset of 500 data points, but stops being very practical as the dataset gets larger. Instead, we want to evaluate each proposal based off a minibatch of data.

Naive Minibatch MH

The simplest way to do this is to just do it: evaluate the loss on a minibatch, keep everything else the same.¹ And this already works ok!^{2.5, 2.75}

But there is an important caveat to consider. This approach is quite sensitive to batch size---for example, it's hard to get good results with a batch size of 1. We can understand why this is by drilling down a bit into what the Metropolis-Hastings algorithm is doing for us.

From a theoretical point of view, the Metropolis-Hastings algorithm is generating samples from a distribution. The distribution we're interested in modeling is the density function $p(\theta\vert \{x_i\})$, the likelihood given to a certain set of model parameters $\theta$ conditioned on the dataset $\{x_i\}$. We compute this using Bayes' rule, and after some simplifying assumptions,³ we get $$ p(\theta\vert\{x_i\}) = \prod_i p(x_i\vert\theta). $$ We can convert this into a loss function by taking negative log,⁴ and for convenience taking the mean of sample losses rather than the sum: $$ \ell(\theta\vert\{x_i\})=-\frac{1}{N}\sum_i\log\big(p(x_i\vert\theta)\big),\qquad\qquad N=\text{number of samples}. $$ So running MH with this full-batch loss function will model the distribution $p(\theta\vert\{x_i\})$.

What happens if we instead compute the loss $$ \ell(\theta\vert B)= -\frac{1}{b}\sum_{x\in B}\log\big(p(x\vert\theta)\big)\qquad\qquad b=\vert B\vert\;\;\text{batch size} $$ on a mini-batch $B\subset\{X_i\}$? It doesn't exactly take a rocket scientist to guess⁵ that instead of getting the full posterior distribution $p(\theta\vert\{x_i\})$ we get something more like a mixture of minibatch posteriors $\sum_B p(\theta\vert B)$ over minibatches $B\subset\{x_i\}$.

How bad is this? Well, in practice these distributions are completely intractible, which is why we're doing MH in the first place, so we'll just have to experiment and see.⁶

Fancypants Minibatch MH

There are a few ways we can try to improve on the naive approach; we'll look at two related ways.

The idea is that instead of taking the minibatch loss at face value, we should treat it as an estimate of the full-batch loss. Write $$ \Delta(\theta', \theta) = \ell(\theta\vert\{x_i\}) - \ell(\theta'\vert\{x_i\}) = \frac{1}{N}\sum_i\log\left(\frac{p(\{x_i\}\vert\theta')}{p(\{x_i\}\vert\theta)}\right) $$ for the log acceptance ratio, the difference in loss between the current parameters $\theta$ and the proposed parameters $\theta'$; and the minibatch version $$ \Delta^*(\theta', \theta) = \ell(\theta\vert B) - \ell(\theta'\vert B) = \frac{1}{b}\sum_{x\in B}\log\left(\frac{p(\{x\}\vert\theta')}{p(\{x\}\vert\theta)}\right). $$ Then $\Delta^*$ is an unbiased estimate of $\Delta$. But the problem is that the same is not true of the acceptance ratios $e^{\Delta^*}$ and $e^{\Delta}$, so the target distribution in MH is not preserved.

Barker's Acceptance Test

To get around this problem we can actually change the acceptance test. Instead of accepting the proposal with probability $e^\Delta$, we can use Barker's acceptance criterion, the logistic function $\frac{1}{1+e^{-\Delta}}$, and MH will still follow the same target distribution.⁷

This doesn't just get rid of the problem, but it opens up a new approach.^7.5 In the acceptance test, sampling a uniform variable $u\sim U(0,1)$ and testing $u<\frac{1}{1+e^{-\Delta}}$ is equivalent to sampling a logistic variable $x\sim X_{\rm log}$ and testing $\Delta+x > 0$, essentially unwrapping the logistic function.⁸

As we said above, $\Delta^*$ is an estimate of $\Delta$---in fact, the central limit theorem tells us what kind of estimate. If we think of the individual log ratios $\log\left(\frac{p(x_i\vert\theta')}{p(x_i\vert\theta)}\right)$ as a distribution, the mean is $\Delta$ and let's call the variance $\sigma^2$. Then $\Delta^*$ is the sample mean on a minibatch, and the central limit theorem states that (for a ``large'' batch size $b$) the sampling distribution of $\Delta^*$ is normally distributed with mean $\Delta$ and variance $\sigma^2/b$.

In other words, $\Delta^*\sim \Delta + N(0, \sigma^2/b)$. Testing $\Delta + N(0, \sigma^2/b) > 0$ actually looks a bit like the acceptance test $\Delta+X_{\rm log} > 0$ we proposed above. And if you squint, maybe it looks the same, and we can just use $\Delta^* > 0$ as our acceptance test.^{9, 10}

This approach is more principled than the naive minibatch approach, but still just an approximation. Since we had to squint a bit, the distribution targeted by this MH test is still not the true posterior $p(\theta\vert\{x_i\})$ that we ideally want.

An ``Exact'' Barker Test

If we want to do a bit more work we can be a bit more exact.¹¹

The approximation we accepted in the previous approach was substituting a normal distribution $N(0, \sigma^2/b)$ for the logistic distribution $X_{\rm log}$. We can improve on this using the insight that $X_{\rm log}$ can actually be written as a sum of our normal distribution and a ``correction'' distribution, $$ X_{\rm log}\sim N(0, \sigma^2/b) + X_{\rm correction}. $$ This isn't just possible for any two distributions, but in this case the conditions are satisfied,¹² at least if we can control the variance $\sigma^2/b$.

Equipped with our correction distribution, we draw a minibatch $B$, compute the sample mean of the logratio $\Delta^*=\frac{1}{b}\sum_{x\in B}\log\left(\frac{p(x\vert\theta')}{p(x\vert\theta)}\right)$ and estimate the variance of the sampling distribution of the mean $\mathrm{Var}(\Delta^*)=\sigma^2/b$ as $\frac{1}{b}\mathrm{Var}\left(\left\{\log\frac{p(x\vert\theta')}{p(x\vert\theta)}\right\}_{x\in B}\right)$.¹³ If the variance is small enough, we can draw from our correction distribution¹⁴ to get $$ \Delta^* + X_{\rm correction}\sim\Delta+N(0,\sigma^2/b)+X_{\rm correction}\sim \Delta+X_{\rm log}, $$ so testing $\Delta^*+X_{\rm correction}>0$ reproduces the Barker acceptance criterion. If the variance is too big, we increase the size of our minibatch and check again. Increasing the batch size will decrease the variance through the factor of $1/b$, which limits how much the minibatch is forced to grow; and the learning rate and temperature will also modulate the baseline variance $\sigma^2$, so can be tuned to admit small minibatch sizes.^15.5

With this adaptive minibatch size and correction variable, we have a minibatch MH test that targets the exact posterior distribution $p(\theta\vert\{x_i\})$.¹⁶

Does it work tho

Time to put these new optimizers to the test! There are two main questions that we want to investigate for each of our new minibatch optimizers.

1. Does using minibatch speed up training?
2. Does using minibatch degrade performance?

So we'll evaluate each model on overall performance, training ``speed'' measured by data consumption, and training speed measured by compute time. Since we're interested in the effect of batch size, we'll try each minibatch model both with a ``regular'' batch size of 64, and test the limits with a batch size of 2.

For these tests we'll use the UCI handwritten digits dataset (much like MNIST, but 8x8 rather than 28x28).

Still not a huge dataset (some 5K samples), nor a terribly difficult one. It's a good first test, but we'll likely see stronger differentiation later with a more challenging dataset.

Despite the image format, we'll just use linear layers, which should still get good results on such a simple task.

Sequential model
  Linear: weight (64, 32), bias (32,)
  ReLU
  Linear: weight (32, 32), bias (32,)
  ReLU
  Linear: weight (32, 10), bias (10,)
total parameters: 3,466
        

Overall Performance

First off, let's see about overall performance. Each optimizer was run for 30 seconds and accuracy evaluated on the hold-out validation set.^16.5

full batch     :  92.7% accuracy
Barker b=64    :  90.4% accuracy
naive b=64     :  90.3% accuracy
exact b=64     :  84.6% accuracy
Barker b=2     :  74.7% accuracy
naive b=2      :  70.6% accuracy
exact b=2      :  40.7% accuracy
        

The results are aligned with what we'd expect: full batch MH has a modest performance edge over minibatch MH (on a dataset where full batch is even possible), but minibatch MH stands up quite well---at least when it has a large enough batch size, as the small batch size of 2 resulted in quite a serious performance drop.

Possibly surprising is the fact that the exact minibatch test performed the worst. This is particularly surprising on the small minibatch, since the exact test uses an adaptive minibatch size and therefore should not be restricted by the small batch size the same way the other tests are.

This may reflect more on the temperment of the exact test rather than its potential; in other words, the exact test has more to gain from precise tuning and scheduling than the other tests. One sign of this is that the exact test is far more sensitive to hyperparameters than the others,¹⁷ as illustrated in the heatmaps generated by Optuna in the course of a hyperparameter search.

As these figures show, the exact test has a very narrow range of effective parameters, compared to a much wider range for the naive test. As training progresses and the model moves into a different region of the loss landscape, the range of effective parameters may shift, and training slows down. Naturally, the effect of this would be more severe for a model which is more sensitive to hyperparameters.

For now we're not using any scheduling or adaptivity for the hyperparameters, but it seems likely that it would give the exact model a considerable performance improvement.

Training Efficiency

The results for training are also generally in line with expectations. We'll start with the efficiency measured by data consumption.

Note that the x-axis is a log scale---minibatch optimizers have finished before the full-batch optimizer even gets out of bed. The smaller batch size also results in less data consumed in training, but of course at a higher cost to performance; and the exact minibatch test consumes more data than the other minibatch tests to get going, but didn't use more data overall in this test.¹⁸

The comparison by training time is more competitive.

The main takeaway here is that the naive and Barker minibatch tests were faster, but as said before, hit their ceiling lower than full-batch. The exact minibatch test is somewhat slower in exchange for its stronger theory, but remains to be seen if more careful treatment can have it outcompete the less sophisticated algorithms, as well as potentially speed it up.

On the whole, good fun to work through the theory and implementation here, and promising results worthy of being tested on a bigger badder dataset!¹⁹

← Micropolis I