Advances in few-shot learning: reproducing results in PyTorch | by Oscar Knagg | Towards Data Science

Few-shot learning is an exciting field of machine learning which aims to close the gap between machine and human in the challenging task of learning from few examples. In my previous post I provided a high level summary of three cutting edge papers in few-shot learning — I assume you’ve either read that, are already familiar with these papers or are in the process of reproducing them yourself.
In this post I will guide you through my experience in reproducing the results of these papers on the Omniglot and miniImageNet datasets, including some of the pitfalls and stumbling blocks on the way. Each paper has its own section in which I provide a Github gist with PyTorch code to perform a single parameter update on the model described by the paper. To train the model just have to put that function inside a loop over the training data. Less interesting details such as dataset handling are omitted for brevity.
Reproducibility is very important, it is the foundation of any field that claims to be scientific. This makes me believe that the prevalence of code-sharing and open-sourcing in machine learning is truly admirable. While publishing code alone is not reproducibility (as there may be implementation errors) it opens up researchers methods to public scrutiny and more importantly accelerates the research of others in the field. In light of this I’d like to thank the authors of these papers for sharing their code as well as any others who’ve open-sourced their implementations.
For the full implementation please see my Github repo at https://github.com/oscarknagg/few-shot
There are two image datasets on which few-shot learning algorithms are evaluated. The first is the Omniglot dataset, which contains 20 images each of roughly 1600 characters from 50 alphabets. These images are typically 28x28 grayscale which is one reason why this dataset is often called the transpose of MNIST.
The second is the miniImageNet dataset, a subset of ImageNet intended to be a more challenging benchmark without being as cumbersome as the full ImageNet dataset. miniImageNet consists of 60,000, 84x84 RGB images with 600 images per class.
In both cases the classes in the training and validation sets are disjoint. I did not use the same training and validation splits as the original papers as my goal is not to reproduce them down to their last minute detail.
In Matching Networks Vinyals et al introduce the idea of a fully differentiable nearest neighbours classifier that is both trained and tested on few-shot tasks.
The Matching Networks algorithm can be summarised as follows:
First embed all samples (query and support set) using an encoder network (4 layer CNN in this case). This is performed by model.encode() (line 41).Optionally calculate full context embeddings (FCE). An LSTM takes the original embeddings as inputs and outputs modified embeddings, taking into account the support set. This is performed by model.f() and model.g() (lines 62 and 67).Calculate pairwise distances between query samples and support sets and normalise using softmax (lines 69 to 77)Calculate predictions by taking the weighted average of the support set labels with the normalised distance (lines 83–89)
First embed all samples (query and support set) using an encoder network (4 layer CNN in this case). This is performed by model.encode() (line 41).
Optionally calculate full context embeddings (FCE). An LSTM takes the original embeddings as inputs and outputs modified embeddings, taking into account the support set. This is performed by model.f() and model.g() (lines 62 and 67).
Calculate pairwise distances between query samples and support sets and normalise using softmax (lines 69 to 77)
Calculate predictions by taking the weighted average of the support set labels with the normalised distance (lines 83–89)
Some things to note:
In this example the x Tensor contains first the support set samples and then the query. For Omniglot it will have shape (n_support + n_query, 1, 28, 28)
The math in the previous post is for one query sample but Matching Networks are in fact trained with a batch of query samples of size q_queries * k_way
I was unable to reproduce the results of this paper using cosine distance but was successful when using l2 distance. I believe this is because cosine distance is bounded between -1 and 1 which then limits the amount that the attention function (a(x^, x_i) below) can point to a particular sample in the support set. Since cosine distance is bounded a(x^, x_i) will never be close to 1! In the case of 5-way classification the maximum possible value of a(x^, x_i) is exp(1)/ (exp(1) + 4*exp(-1)) ≈ 0.65. This led to very slow convergence when using cosine distance.
I think it’s possible to reproduce the results using cosine distance with either longer training times, better hyperparameters or a heuristic like multiplying the cosine distance by a constant factor. Seeing as the choice of distance is not key to the paper and results are very good using l2 distance I decided to spare myself that debugging effort.
In Prototypical Networks Snell et al use a compelling inductive bias motivated by the theory of Bregman divergences to achieve impressive few-shot performance.
The Prototypical Network algorithm can be summarised as follows:
Embed all query and support samples (line 36)Calculate class prototypes taking the mean of the embeddings of each class (line 48)Predictions are a softmax over the distances between the query samples and the class prototypes (line 63)
Embed all query and support samples (line 36)
Calculate class prototypes taking the mean of the embeddings of each class (line 48)
Predictions are a softmax over the distances between the query samples and the class prototypes (line 63)
I found this paper delightfully easy to reproduce as the authors provided the full set of hyperparameters. Hence I was easily able to achieve the stated performance to within ~0.2% on the Omniglot benchmark and within a few % on the miniImageNet benchmark without having to perform any tuning of my own.
In MAML Finn et al introduce a powerful and broadly applicable meta-learning algorithm to learn a network initialisation that can quickly adapt to new tasks. This paper was the most difficult yet most rewarding to reproduce of the three in this article.
The MAML algorithm can be summarised as follows:
For each n-shot task in a meta-batch of tasks, create a new model using the weights of the base model AKA meta-learner (line 79)Update the weights of the new model using the loss from the samples in the task by stochastic gradient descent (lines 81–92)Calculate loss of the updated model on some more data from the same task (lines 94–97)If performing 1st order MAML update the meta-learner weights with the gradient of the loss from part 3. If performing 2nd order MAML calculate the derivative of this loss with respect to the original weights (lines 110+)
For each n-shot task in a meta-batch of tasks, create a new model using the weights of the base model AKA meta-learner (line 79)
Update the weights of the new model using the loss from the samples in the task by stochastic gradient descent (lines 81–92)
Calculate loss of the updated model on some more data from the same task (lines 94–97)
If performing 1st order MAML update the meta-learner weights with the gradient of the loss from part 3. If performing 2nd order MAML calculate the derivative of this loss with respect to the original weights (lines 110+)
The biggest appeal of PyTorch is its autograd system. This is a piece of code-magic that records operations acting on torch.Tensor objects and dynamically builds the directed acyclic graph of these operations under the hood. Backpropagation is as simple as calling .backwards() on the final result. I had to learn a bit more about this system in order to calculate and apply parameter updates to the meta-learner, which I will now share with you.
Typically when training a model in PyTorch you create an Optimizer object tied to the parameters of a particular model.
from torch.optim import Adamopt = Adam(model.parameters(), lr=0.001)
When opt.step() is called the optimiser reads the gradients on the model parameters and calculates an update to those parameters. However in 1st order MAML we’ re going to calculate the gradients using one model (the fast weights) and apply the update to a different model i.e. the meta-learner.
A solution to this is to use an under-utilised bit of PyTorch functionality in the form of torch.Tensor.register_hook(hook). Register a hook function to a Tensor and this hook function will be called whenever a gradient with respect to this tensor is computed. For each parameter Tensor in the meta-learner I register a hook that simply replaces the gradient with the corresponding gradient on the fast weights (lines 111–129 in gist). This means that when opt.step() is called the gradients of the fast model will be used to update the meta-learner weights as desired.
When making my first attempt at implementing MAML I instantiated a new model object (subclass of torch.nn.Module) and set the values of its weights equal to the meta-learner’s weights. However this makes it impossible to perform 2nd order MAML as the weights of the fast model are disconnected from the weights of the meta-learner in the eyes of torch.autograd. What this means is when I call optimiser.step() (line 140 in the gist) the autograd graph for the meta-learner weights is empty and no update is performed.
# This didn't work, meta_learner weights remain unchangedmeta_learner = ModelClass()opt = Adam(meta_learner.parameters(), lr=0.001)task_losses = []for x, y in meta_batch:    fast_model = ModelClass()    # torch.autograd loses reference here!    copy_weights(from=meta_learner, to=fast_model)    task_losses.append(update_with_batch(x, y))meta_batch_loss = torch.stack(task_losses).mean()meta_batch_loss.backward()opt.step()
The solution to this is functional_forward() (line 17) which is a slightly awkward hack that manually performs the same operations (convolution, max pooling, etc...) as the model class using torch.nn.functional. This also means that I have to manually perform a parameter update of the fast model. The consequence of this is that torch.autograd knows to backpropagate gradients to the original weights of the meta-learner. This leads to a spectacularly large autograd graph.
However 2nd order MAML is a trickier beast than just that. When I first wrote my 2nd order MAML implementation I thought I had got everything to work miraculously on the first try. At least there were no exceptions right? Only after running a full set of Omniglot and miniImageNet experiments did I begin to doubt my work — the results were just too similar to 1st order MAML. This is typical of an unfortunate breed of silent ML bugs which don’t cause exceptions but only become visible in the final performance of a model.
Hence I decided to buckle down and write a unit test that would confirm that I was truly performing a 2nd order update. Disclaimer: in the spirit of true test-driven development I should’ve written this test before running any experiments. It turns out I’m not perfect 😛.
The test I decided on was to run the meta_gradient_step function on a dummy model and manually parse the autograd graph, counting the number of double backwards operations. This way I can be absolutely sure that I am performing a 2nd order update when desired. Conversely, I was able to test that my 1st order MAML implementation only performs a 1st order update with no double backwards operations.
I finally located the bug to not applying the create_graph parameter in the inner training loop (line 86). I was, however, retaining the autograd graph of the loss on the query samples (line 97) but this was insufficient to perform a 2nd order update as the unrolled training graph was not created.
Training time was quite long (over 24 hours for the 5-way, 5-shot miniImageNet experiment) but in the end I had fairly good success reproducing results.
I hope that you’ve learnt something useful from this technical deep dive. If you have any questions feel free to let me know in the comments.
If you want to dig in further check out the full implementation at https://github.com/oscarknagg/few-shot