SKM Embedding of MNIST

I recently thought up a machine learning algorithm called "Smooth Kernel Macines" (SKM). In this post I will try out SKM on the ever ubiquitous MNIST dataset. The goal of this post is not so much to achieve state of the art performance on MNIST (though that would be nice), as it is to simply try out SKM on a familiar and well understood dataset.

tldr; I achieve a respectable 0.006 error rate using an SKM type layer on top of a convolutional neural net feature extractor. An SKM output layer works a little better than a K way softmax (at least for MNIST). SKM trains faster, and comes with an accurate built in measure of prediction confidence.

Data Loading and Preprocessing¶

We need to load up the data, convert the digits to a one-hot type binary encoding so that averages of the label vectors make sense and then normalize the digit images to have unit norm and zero mean.

Feature Extractor¶

To extract the features for the SKM model I have decided to use a three sets of of 3x3 convolution layers feeding into 1x1 Squeeze Layers feeding in to max pooling layers with stride 2. I have left out the max pooling layer on the final set of 3x3 and 1x1 convolutions letting the dense output layers do the remaining dimensionality reduction. I have opted for exponential linear units (ELU) instead of the more often employed rectified linear units (ReLU) out of an intuition that the sharp edges of the ReLU will cause more spurious similarities since the distance between any 2 points on the saturated part of the ReLU will be 0 but the distances go smoothly towards 0 in the case of the ELU. However the hand full of numerical experiments I have done seem to suggest that although the ELU may in fact perform slightly better there is not a great deal of difference in performance at least for this dataset.

In order to be able to exactly clone the weights from the SKM network into a regular softmax output network for comparison purposes I have opted for a 10 dimensional similarity space. This turns out to be slightly more dimensions than is necessary for the similarity embedding (see the PCA analysis below) but the hand full of extra dimensions don't seem to adversely affect model performance.

In order to be sure that the difference in performance of the softmax network and the performance of the SKM network is not just due to a lucky (or unlucky) set of starting weights we clone the starting weights of one network into the other so that we can meaningfully compare their initial performance and training curves.

At just around 200,000 total parameters the model is pretty light weight. Many fewer parameters than is typically used for example at the time of this writing the current TensorFlow MNIST tutorial architecture uses over 3 million parameters and achieves a test accuracy of 97.3% after 33 training epochs. This seems like an intentionally weak choice of network architecture (maybe to allow people the thrill of improving the networks performance by making tweaks?) By contrast using a 3x3 instead of 5x5 kernel size and a smaller intermediate dense layers and slightly heavier dropout regularization the keras MNIST example uses just around 2 million parameters and hits 99.25% test set accuracy.

It is difficult to count up parameters in papers in the literature but as best as I can tell looking at the papers with close to state of the art performance found here I think parameter counts of a couple million is typical. By fiddling with the architecture I have managed to get as high as 99.2 test accuracy using only 70,000 parameters but at the end of the day performance matters more than network size, a net that has a 50% lower error rate at the cost of 100% more computation is almost always going to be worth the extra overhead. This lightweight 200,000 parameter network works almost as well as any larger network I tried out and so we will stick with it.

Kernel Function¶

In the first SKM post we used a Gaussian kernel similarity function and there was no need to fiddle with the width parameter of the Gaussian. However in the 10 dimensional similarity space that we are using in this post the ability of the model to productively train becomes very sensitive to the choice of kernel width parameter of a Gaussian kernel. If the kernel is too narrow then all points are effectively isolated which leads to either extremely inefficient learning or complete model divergence (usually the latter). If the similarity kernel is too wide then the similarity of the closest points are all essentially the same. The points which most strongly drive the gradient are those at distances which are very near to the kernel width $\sigma$ points much closer than that and points much farther away both having little to no effect. This leads to non-intuitive effects and very inefficient model training or outright model divergence (usually the former).

Ideally we would like to avoid the need to carefully tune the kernel width of our models as a hyper parameter. To avoid the first trouble of an overly narrow kernel leading to isolated points with effectively 0 similarity to each other we just need to make the kernel fall off a little more slowly. To handle the second problem of indifference of the gradient to points near r=0 we simply need to pick a similarity kernel such that not only does the value of the similarity fall off to 0 with increasing distance but so does the first derivative.

The first function which comes to mind as an alternative to the Gaussian kernel is the exponential kernel. The tails of the exponential kernel fall off much more slowly than the Gaussian meaning that we have to worry less about making the similarity kernel too narrow. The derivative of the exponential similarity with respect to distance is proportional to the similarity itself, thus the training gradient is coupled to pairs of points directly in proportion to the degree of similarity between those points which seems ideal. In fact in this experiment the choice of exponential kernel seems to serve us well.

There is a slight complication due to the fact that the limited numerical precision the pairwise squared distance that comes out of the tf_pairwise_sq_distance_matrix function is sometimes very slightly negative. To solve this we have added a small positive constant to all squared distances prior to taking the square root to ensure positivity.

Unfortunately the use of the exponential kernel instead of the Gaussian kernel doesn't completely eliminate the need for tuning the kernel width as a hyper parameter but the few experiments I have carried out show it to be much more forgiving than the Gaussian kernel.

Stratified Training¶

With 10 different classes the amount of variation batch to batch in the number of included examples of each class differs significantly. Since the SKM model relys entirely on the apparent similarity of input data it helps with smooth training to always pass in equal numbers of examples of each digit in every batch. There aren't exactly the same number of examples of each category so we end up cutting out a small number of samples each epoch. But the returns in terms of training stability and end model quality are worth it.

Randomly Sampled vs. Nearest Neighbor Prediction¶

Since the train and test datasets are too large to be able to comfortably hold the entire train to test similarity matrix in memory all at once we can either use a nearest neighbor search to cherry pick only those training samples which are closest to a given test sample or we can pick a few random subsets of the training data and use those as prototypes.

The sampling approach is appealing because it directly applies the model we are actually training. However, I would naievely expect the nearest neighbor approach to work a little better than the sampling method since we are guaranteed to find the most salient examples each time. But it is difficult to know how much of the nearest neighbor models success is due to the nearest neighbor algorithm itself versus that due to the learning occuring in the feature extraction network. We will evalute the model both via a direct sampling method and a nearest neighbor method at each epoch and compare the results.

Performance Comparison¶

The error rate should be treated as a stochastic number and not as an exact value. As a good general rule of thumb the 1 sigma error on any integer count is simply the square root of the count itself (which is where we are getting the errorbars in the above plot). Although this approximation breaks down for very small counts it is still an excellent approximation in this regime where whe have around 60 misclassified examples. Thus the 0.002 difference between the SKM error rate and the softmax error rate should be seen as being uncertain at the level of at least around $\sqrt{60}/10000 \approx 0.0008$ for the regular softmax model. Seen in that light the improvement of 0.002 represents just around a 2.5 sigma significance level, which is suggestive but nothing to hang your hat on.

You might be suspicious of that error calculation. Shouldn't we add the uncertainty in both error rates in quadrature? That would be a good approximtion if the two errors were uncorrelated with each other. But since we are evaluating the models on the very same data and even have gone out of our way to train the models in the same way and with the same architecture the two error rates can hardly be considered to be independent of each other. In fact the correlation between the two may well be high enough that the difference between the two error rates might actually have a smaller effective uncertainty than either one individually but there is no way to know apriori how highly correlated we should treat the two numbers.

The SKM model learns much faster than the softmax network. The SKM model hits 98.7% accuracy after just one training epoch a full percent higher than that of the regular softmax model. At least some of this rapid training in early iterations is due to the fact that even without any training at all the similarity between pairs of digits of the same class is slightly larger than the similarities between different classes. Thus the SKM model always starts from a model which slightly beats random chance predictions whereas the regular softmax network will usually start from a model which is slightly worse than random chance.

Although the sampled SKM model somewhat underperforms relative to the nearest neighbor version at first it eventually comes to very nearly match its performance in later epochs. The nearest neighbor search is actually many times faster to carry out than the sampling estimator and so the nearest neighbor estimator is the clear winner. However the sampling estimator retains value as a diagnostic tool especially since it more closely mimics what the model looks like from the perspective of a single training batch. If the sampling estimator doesn't do a good job that may mean individual training batches don't provide enough information to drive meaningful gradients.

Ideally it would be nice to see a series of experiments on lots of different data sets and a proper comparison of the performance of the networks when each network has its hyper-parameters set via search to be optimized for their respective models instead of sharing the hyper parameter values as we have done here to avoid hyper parameter searching. I have made some preliminary SKM models on several different more challenging data sets such as fashion MNIST, cifar10, cifar100, and SVHN and may do a comparison post in the future so stay tuned.

Embedding Dimensionality¶

The use of a full 10 dimensions for the SKM embedding is probably sub-obpimal there is no need for all that space for this particular problem. Because the SKM relys on close collisions of prototype examples with the target examples for prediction it helps to keep the dimensionality of the similarity space into which we embed points relatively small. The larger the dimension the less likely two points are to get mapped near to each other.

During training the advantages of having a lower effective dimension outweigh the potential extra freedom of having more space to map points to and so the model learns to restrict itself to a low dimensional subspace of the embedding. For this data and choice of kernel it seems that the model benefits from having perhaps 6 dimensions and learns to effectively restrict itself onto a 6 dimensional subspace. You may suppose that the apparent reduction in the number of dimensions is a property of the PCA and not a property of the data or the model but I can tell you that even if you use a 20 or 30 dimensional similarity space the model still learns to restrict itself to a subspace of effectively the same dimension with almost no variance at all in the final dozen or more dimensions.

Prediction Confidence¶

Although the performance difference in the models is small and may even be spurious where the SKM model really shines is in the fact that it comes with a built in confidence mechanism. Ideally our class probabilities would be well calibrated so that an output of probability 0.9 would actually mean that we can expect the predicted class to be right 90% of the time and wrong 10% of the time. Unfortunately coming up with well calibrated probabilities (and being confident that they are in fact well calibrated) is actually a MUCH harder problem than coming up with reasonable probability estimates.

For example take the classic probability problem of randomly drawing balls out of a bucket (with replacement). If you draw 1 ball out and it is red you might assign a probability of 100% to the next ball being red. If you are a little more reasonable (or are a bayesian instead of a frequentist) you would perhaps hedge your bets a little and apply a weak prior. With a typical bit of baysian logic you might add a "pseudocount" of 1/2 to your observations of red and not-red balls. With those numbers you would come up with a probability estimate (1+1/2)/(1+1/2+1/2) = 0.75 which seems reasonable to me.

If you were offered an even odds bet that the next ball is some other color than red you probably shouldn't take the bet (or at any rate I would not take that bet). Just about any reasonable (weak) prior that you might pick will lead you to an estimate of the probability that is greater than 50%. However, I think it is intuitively clear that the confidence we should assign to the chances that the true probability is greater than 50% should not be high. The very fact that the bet is being offered at all may mean that whoever is offering the bet knows something we don't, and we should be wary.
Contrast that to the situation where you have drawn 1,000 samples and have drawn 75% reds. Although your probability estimate is the same as before your confidence in that probability is vastly greater and now you can take the bet confidently. The probability of the true red fraction being below 50% in that case is vanishingly small.

To swing back around to the problem of predicting a digit label given an image. Instead of drawing balls from a bucket we are drawing test images from our hold out set and instead of "red" and "not red" we are recording correct and incorrect predictions. If we put all 10,000 of our test images into one "bucket" so to speak then we end up with the global error estimate (and uncertainties in that error) that we used in the test accuracy vs epoch plot above.

However not all of our predictions are of the same quality and we needn't lump them all together when we figure out how well our model performs. For example if our output probabilities are well calibrated then the predictions of digits with a maximum class probability of 90% will be right 90% of the time and are high quality predictions. On the other hand if the probabilities are more spread out so that the maximum probability of any one class is just 60% that prediction is of a somewhat lower quality.

In the case of our softmax model if a new training sample explores a region of activation space which is extremely different from anything ever seen by the network during training then the resultant prediction is likely to be nonsense. In these edge cases our model may give us a predicted class likelihood for a certain class of 99.999% and still be totally wrong. Because the softmax network retains no information about what "typical" behavior looks like it very happily will return a very confident prediction even for extremely odd examples.

On the other hand SKM works entirely by making a sort of analogy between the current example and the set of prototype examples available to it. Thus if the SKM model encounters a sample which explores a completely new region of activation space never before seen it will be far from all known examples and get a very low total similarity sum. The similarity sum is a sort of estimate of the density of training prototypes in the space near the target image to be predicted. If the density of prototypes near an image is low then that image represents something new and unusal and we should be suspicious of any predictions we make. On the other hand if there are lots of nearby example prototypes then we are on well trodden metaphorical ground and the predictions of the SKM model should be of high quality.

Lets see how well the probability of the softmax model and/or the neighbor count of the SKM model predicts the quality of our answers and compare the two.

Although both the softmax model and SKM model get more than 99% of the digit labels right the softmax model gives very little indication of when its answers might possibly be wrong. Just about 50% of all of the 80 incorrectly labeled digits get right up to the maximum allowed probability for their (incorrect) predicted label. On the other hand the SKM model only makes incorrect predictions when the digits in question get mapped to very unusual places in the embedding space, as indicated by the relatively low similarities of the nearest points.

Having a good way to know when your model's predictions are likely to be poor is just as valuable as knowing how well your model is likely to perform. Especially in the cases where models are driving decisions it can be incredibly valuable to know whether that 99% class prediction represents a well considered 99% drawing from thousands of training examples or the vacuous sort of 99% class prediction that arose because we have never seen anything like that situation before.

Suppose we feed an image through the softmax classifier and get a maximum predicted label probability of 90%. Further suppose we are offered a bet wherein we will pay out 2 dollars if the prediction turns out to be incorrect and will be paid 1 dollar if the prediction is incorrect. You may naively assume that the expectation of that bet is 70 cents. Since the softmax model is right 99.2 percent of the time it must be pretty good and if we trust its prediction of a 90% chance of winning a dollar and a ten percent chance of losing 2 dollars we come out with a 70 cent gain on average.

But if we look at the plot above in the few cases where no one class got more than a 90% predicted probability we see a very different story. The softmax model is almost always essentially 100% certain of every category and in the very few cases where the model doesn't come out with a probability of 99.9% for just one particular label then the numbers of correctly and incorrectly classified digit examples is actually roughly half and half. So the information that our model was only 90% certain actually means that we should expect to lose 2 dollars half the time and win 1 dollar half the time making the expectation of that bet about negative one dollars. In fact because the most common probability assigned to incorrect predictions by the softmax model is near 100% there are not really any predictions we can be more confident of than the general prediction.

On the other hand if we consider the density of prototype data via the SKM model we can clearly see that if we are granted the information that a test image has a prototype density at least as high as the average density then the error rate is so low that the 10,000 held out test images is just not enough to estimate it. That is pretty cool.

Embedding Space Visualization¶

Lets plot up 2 of the 10 dimensions of the embedding space and see how the densities of different digits lay out.

Image Distance Distributions¶

Finally because the whole SKM classification process depends on inter image distances lets plot up some distance and similarity matrices and see how the two tend to interact.

It is clear to see that the learned embedding has indeed managed to locate images belonging to the same digit much more closely together than those belonging to different digits. From the perspective of maximizing the probability of the target class it would be ideal to have the similarity of all non-similar digits be exactly 0. However due to our slowly decaying exponential kernel nearly all images have a non-zero similarity to each other. For the reasons that we discussed in the kernel section above this actually leads to more robust training of the model. Though what the choice of optimal kernel might be is very much an open question both in terms of functional form and/or any associated hyper parameters.