M EMORY - BASED P ARAMETER A DAPTATION
A BSTRACT
Deep neural networks have excelled on a wide range of problems, from vision to language and game playing.
Neural networks very gradually incorporate information into weights as they process data, requiring very low learning rates.
If the training distribution shifts, the network is slow to adapt, and when it does adapt, it typically performs badly on the training distribution before the shift.
Our method,
Memory-based Parameter Adaptation, stores examples in memory and then uses a context-based lookup to directly modify the weights of a neural network.
Much higher learning rates can be used for this local adaptation, reneging the need for many iterations over similar data before good predictions can be made.
As our method is memory-based, it alleviates several shortcomings of neural networks, such as catastrophic forgetting, fast, stable acquisition of new knowledge, learning with an imbalanced class labels, and fast learning during evaluation.
We demonstrate this on a range of supervised tasks: large-scale image classiﬁcation and language modelling.
I NTRODUCTION
Neural networks have been proven to be powerful function approximators, as shown in a long list of successful applications: image classiﬁcation (e.g. Krizhevsky et al., 2012), audio processing (e.g.
Oord et al., 2016), game playing (e.g. Mnih et al., 2015; Silver et al., 2017), and machine translation (e.g. Wu et al., 2016).
Typically these applications apply batch training to large or near-inﬁnite data sets, requiring many iterations to obtain satisfactory performance.
Humans and animals are able to incorporate new knowledge quickly from single examples, continually throughout much of their lifetime.
In contrast, neural network-based models rely on the data distribution being stationary and the training procedure using low learning rates and many passes through the training data to obtain good generalisation.
This limits their application to life-long learning or dynamic environments and tasks.
Problems in continual learning with neural networks commonly manifest as the phenomenon of catastrophic forgetting (McCloskey & Cohen, 1989; French, 1999): a neural network performs badly on old tasks having been trained to perform well on a new task.
Several recent approaches have proven promising at overcoming this, such as elastic weight consolidation (Kirkpatrick et al., 2017).
Recent work in language modelling has demonstrated how popular neural language models may appropriately be adapted to take advantage of rare, recently seen words, as in the neural cache (Grave et al., 2016), pointer sentinel networks (Merity et al., 2016) and learning to remember rare events (Kaiser et al., 2017).
Our work generalises these approaches and we present experimental results where we apply our model to both continual or incremental learning tasks, as well as language modelling.
We propose Memory-based Parameter Adaptation (MbPA), a method for augmenting neural networks with an episodic memory to allow for rapid acquisition of new knowledge while preserving
[Figure]
the high performance and good generalisation of standard deep models.
It combines desirable properties of many existing few-shot, continual learning and language models.
We draw inspiration from the theory of complementary learning systems (CLS: McClelland et al., 1995; Leibo et al., 2015;
Kumaran et al., 2016), where effective continual, life-long learning necessitates two complementary systems: one that allows for the gradual acquisition of structured knowledge, and another that allows rapid learning of the speciﬁcs of individual experiences.
As such, MbPA consists of two components: a parametric component (a standard neural network) and a non-parametric component (a neural network augmented with a memory containing previous problem instances).
The parametric component learns slowly but generalises well, whereas the non-parametric component rapidly adapts the weights of the parametric component.
The non-parametric, instance-based adaptation of the weights is local, in the sense the modiﬁcation is directly dictated by the inputs to the parametric component.
The local adaptation is discarded after the model produces its output, avoiding long term consequences of strong local adaptation (such as overﬁtting), allowing the weights of the parametric model to learn slowly leading to strong performance and generalisation.
The contributions of our work are: (i) proposing an architecture for enhancing powerful parametric models with a fast adaptation mechanism to efﬁciently cope with changes in the task at hand; (ii) establish connections between our method and attention mechanisms frequently used for querying memories; (iii) present a Bayesian interpretation of the method allowing a principled form of regularisation; (iv) evaluating the method on a range of different tasks: continual learning, incremental learning and data distribution shifts, obtaining promising results.
M ODEL - BASED P ARAMETER A DAPTATION
Our models consist of three components: an embedding network, f γ , a memory M and an output network g θ .
The embedding network, f γ , and the output network, g θ , are standard parametric (feed forward or recurrent) neural networks with parameters γ and θ , respectively.
The memory M is a dynamically-sized memory module that stores key and value pairs, M = { ( h i , v i ) } .
Keys { h i } are given by the embedding network.
The values { v i } correspond to the desired output y i .
For classiﬁcation, y i would simply be the true class label, whereas for regression, y i would be the true regression target.
Hence, upon observing the j -th example, we append the pair ( h j , v j ) to the memory M , where:
[Equation]
The memory has a ﬁxed size and acts as a circular buffer: when it is full, the oldest data is overwritten ﬁrst.
Retrieval from the memory M uses K -nearest neighbour search on the keys { h i } with
Euclidean distance to obtain the K most similar keys and associated values.
Our model is used differently in the training and testing phases.
During training, for a given input x , we parametrise the conditional likelihood with a deep neural network given by the composition procedure M B PA-T RAIN
Sample mini-batch of training examples B = { ( x b , y b ) } b from training data.
Calculate the embedded mini-batch B (cid:48)
=
{ ( f γ ( x b ) , y b ) : x b , y b ∈ B } .
Update θ , γ by maximising the likelihood (1) of θ and γ with respect to mini-batch B
Add the embedded mini-batch examples B (cid:48) to memory M : M ← M ∪ B (cid:48) .
procedure M B PA-T EST (test input: x , output prediction: ˆ y )
Calculate embedding q = f γ ( x ) , and ∆ total ← 0 .
Retrieve K -nearest neighbours to q and producing context, C for each step of MbPA do
Calculate ∆ M ( x,
+ ∆ according to (
) ∆ total ← ∆
∆
Output prediction
+ total ( h ) of the embedding and output networks.
Namely,
[Equation]
In the case of classiﬁcation, the last layer of g θ is a softmax layer.
The parameters { θ, γ } are estimated by maximum likelihood estimation.
The memory is updated with new entries, as they are seen, however no local adaptation is performed on the model.
Figure 1 (left) shows a diagram of the training setting and Algorithm 1 (MbPA-Train) shows the algorithm for updating MbPA during training.
On the other hand, at test time, it temporarily adapts the parameters of the output network based upon the current input and the contents of the memory M .
That is, it uses the exact same parametrisation as (1), but with a different set of parameters in the output network.
Let the context C of an input x be the keys, values and associated weights of the K nearest neighbours to query q = f γ ( x ) in the memory M : C
= { ( h ( x ) k , v ( x ) k , w ( x ) k ) }
Kk =1 .
The coefﬁcients w ( x ) k
∝ kern ( h ( x ) k , q ) are weightings of each of the retrieved neighbours according to their closeness to the query f γ ( x ) .
kern ( h, q ) is a kernel function which, following (Pritzel et al., 2017), we choose as kern ( h, q )
= 1 (cid:15) + (cid:107) h − q (cid:107) 22 .
The parametrisation of the likelihood takes the form,
[Equation]
as opposed to the standard parametric approach g θ ( f γ ( x )) , where θ x = θ + ∆ M ( x, θ ) with ∆ M ( x, θ ) being a contextual (it is based upon the input x ) update of the parameters of the output network.
The MbPA adaptation corresponds to decreasing the weighted average negative loglikelihood over the retrieved neighbours in C .
Figure 1 (right) shows a diagram of the testing setting and Algorithm 1(MbPA-Test) shows the algorithm for using MbPA during testing.
An interesting property of the model is that the correction ∆ M ( x, θ ) is such that, as the parametric model becomes better at ﬁtting the training data (and consequently the episodic memories), it selfregulates and diminishes.
In the CLS theory, this process is referred to as consolidation, when the parametric model can reliably perform predictions without relying on episodic memories.
2. M AXIMUM A P OSTERIORI I NTERPRETATION OF M B PA
We can now derive ∆ M ( x, θ ) , motivated by considering the posterior distribution on the parameters θ x .
Let x correspond to the input with context C
= { h k , v k , w ( x ) k } Kk =1 .
The maximum a posteriori over the context C , given the parameters obtained after training θ , can be written as:
[Equation]
where the second term is a weighted likelihood of the data in C and log p ( θ x | θ
[Equation]
∝ − ||
[Equation]
(i.e. a
Gaussian prior on θ x centred at θ ) can be thought as a regularisation term that prevents overﬁtting.
See Appendix D for details of this derivation.
[Figure]
Equation (3) does not have a closed form solution, and requires ﬁtting a large number of parameters at inference time.
This can be costly and susceptible to overﬁtting.
We can avoid this problem by adapting the reference parameters θ .
Speciﬁcally, we perform a ﬁxed number of gradient descent steps to minimise (3).
One step of gradient descent to the loss in (3) with respect to θ x yields
[Equation]
where β is a scalar hyper-parameter.
These adapted parameters are used for output computation but discarded thereafter, as described in Algorithm 1.
2. F ROM ATTENTION TO LOCAL FITTING
A standard formulation of memory augmented networks is in the form of attention (Bahdanau et al., 2014), i.e. query memory to use a weighted average based on some similarity metric.
We can now show that an attention-based procedure is a particular case of local adaptation or MbPA.
The details of this are discussed in Appendix E. Effectively, attention can be viewed as ﬁtting a constant function the neighbourhood of memories, whereas MbPA generalises to ﬁt a function parameterised by the output network of our model.
The diagram in Figure 2 illustrates an example in a regression task for simplicity.
Given a query (blue), the retrieved memories and their corresponding values are depicted in red.
The predictions of an attention based model are shown in orange.
We can see that the prediction is biased towards the value of the neighbours with higher functional value.
In magenta we represent the predictions made by the model g θ .
We can see that it is not able to explain all memories equally well.
This could be either because the problem is too difﬁcult, poor training, or because the a prediction needs to be made while assimilating new information.
The green curve show the prediction obtained after adapting the parameters to better explain the episodic memories.
R ELATED WORK
A key component of MbPA is the non-parametric, episodic memory.
Many recent works have looked at augmenting neural network systems with memories to allow for fast adaptation or incorporation of new knowledge.
Variants of this architecture have been successfully used in the context of classiﬁcation (Vinyals et al., 2016;
Santoro et al., 2016; Kaiser et al., 2017), language modelling (Merity et al., 2016; Grave et al., 2016), reinforcement learning (Blundell et al., 2016; Pritzel et al., 2017), machine translation (Bahdanau et al., 2014), and question answering (Weston et al., 2014), to name a few.
For the MbPA experiments below, we use a memory architecture similar to the Differentiable
Neural Dictionary (DND) used in Neural Episodic Control (NEC)
(Pritzel et al., 2017).
One key difference is that we do not train the embedding network through the gradients from the memories (as they are not used at training time).
While many of these approaches share a contextual memory lookup system, MbPA is distinct in the method by which the memories are used.
Matching Networks (Vinyals et al., 2016) use a non- parametric network to map from a few examples to a target class via a kernel weighted average.
Prototypical Networks (Snell et al., 2017) extend this and use a linear model instead of a nearest neighbour method.
MbPA is further related to meta-learning approaches for few shot learning.
In the context of learning invariant representations for object recognition, Anselmi et al.
(2014) proposed a method that can invariantly and discriminatively represent objects using a single sample, even of a new class.
In their method, instead of training via gradient descent, image templates are stored in the weights of simple-complex cell networks while objects undergo transformations.
Optimisation as a model of few shot learning (Ravi & Larochelle, 2016) proposes using a meta-learner LSTM to control the gradient updates of another network, while Model-Agnostic Meta-Learning (MAML Finn et al.
(2017)) proposes a way of doing meta-learning over a distribution of tasks.
These methods extend the classic ﬁne-tuning technique used in domain adaptation type of ideas (e.g. ﬁt a given neural network to a small set of new data).
The MAML algorithm (particularly related to our work) aims at learning an easily adaptable set of weights, such that given a small amount of training data for a given task following the training distribution, the ﬁne-tuning procedure would effectively adapt the weights to this particular task.
Their work does not use any memory or per-example adaptation and is not based on a continual (life-long) learning setting.
In contrast, our work, aims at augmenting a powerful neural network with a ﬁne-tuning procedure that is used at inference only.
The idea is to enhance the performance of the parametric model while maintaining its full training.
Recent approaches to addressing the continual learning problem have included elastic weight consolidation (Kirkpatrick et al., 2017), where a penalty term is added to the loss for deviations far from previous weights, and learning without forgetting (Li & Hoiem, 2016; Furlanello et al., 2016), where distillation (Hinton et al., 2015) from previously trained models is used to keep old knowledge available.
Gradient Episodic Memory for Continual Learning (Lopez-Paz & Ranzato, 2017) attempts to solve the problem by storing data from previous tasks and taking gradient updates when learning new tasks that do not increase the training loss on examples stored in memory.
There has been recent work in applying attention to quickly adapt a subset of fast weights (Ba et al., 2016).
A number of recent works in language modelling have augmented prediction with attention over recent examples to account for the distributional shift between training and testing settings.
Works in this direction include neural cache (Grave et al., 2016) and pointer sentinel networks (Merity et al., 2016).
Learning to remember rare events (Kaiser et al., 2017) augments an
LSTM with a key-value memory structure, and meta networks (Munkhdalai & Yu, 2017) combines fast weights with regular weights.
Our model shares this ﬂavour of attention and fast weights, while providing a model agnostic memory-based method that applies beyond language modelling.
Works in the context of machine translation relate to MbPA.
Gu et al.
(2017) explore how to incorporate information from memory into the ﬁnal model predictions.
The authors ﬁnd that shallow mixing works best.
We show in this paper that MbPA is another competitive strategy to shallow mixing, and often working better (PTB for language modelling, ImageNet for image classiﬁcation).
The work by Li et al.
(2016) shares the focus on fast-adaptation during inference with our work.
Given a test example, the translation model is ﬁne-tuned by ﬁtting similar sentences from the training set.
MbPA can be viewed as a generalisation of such approach: it relies on an episodic memory (rather than the training set), contextual lookup and similarity based weighting scheme to ﬁne-tune the original model.
Collectively, these allow MbPA to be a powerful domain-agnostic algorithm, which allows it to handle continual and incremental learning.
Finally, we mention that our work is closely related to the local regression and adaptive coefﬁcient models literature, see Loader (2006) and references therein.
Locally adaptive methods achieved relatively modest success in high-dimensional classiﬁcation problems, as ﬁtting many parameters to a few neighbours often leads to over ﬁtting.
We attempt to counter this with contextual lookups and a local modiﬁcation of only a subset of model parameters.
E XPERIMENTS AND R ESULTS
Our scheme uniﬁes elements from traditional approaches to continual, one-shot, and incremental or life-long learning.
Models that solve these problems must have certain fundamental attributes in common: the ability to negate the effects of catastrophic forgetting, unbalanced and scarce data, while displaying rapid acquisition of knowledge and good generalisation.
In essence, these problems require the ability to deal with changes and shifts in data distributions.
We demonstrate that MbPA provides a way to address this.
More concretely, due to the robustness of the local adaptation, the model can deal with shifts in domain distribution (e.g. train vs test set in language), the task label set (e.g. incremental learning) or sequential distributional shifts (e.g. continual learning).
Further, MbPA is agnostic to both task domain (e.g. image or language) and choice of underlying parametric model, e.g. convolutional neural networks (LeCun et al., 1998) or
LSTM (Hochreiter & Schmidhuber, 1997).
To this end, our experiments focus on displaying the advantages of MbPA on widely used tasks and datasets, comparing with competing deep learning methods and baselines.
We start by looking at the continual learning framework, followed by incremental learning, the problems of unbalanced data and test time distributional changes.
4. C ONTINUAL L EARNING : SEQUENTIAL DISTRIBUTIONAL SHIFT
In this set of experiments, we explored the effects of MbPA on continual learning, i.e. when dealing with the problem of sequentially learning multiple tasks without the ability to revisit a task.
We considered the permuted MNIST setup (Goodfellow et al., 2013).
In this setting, each task was given by a different random permutation of the pixels of the MNIST dataset.
We explored a chaining of 20 different tasks (20 different permutations) trained sequentially.
The model was tested on all tasks it had been trained on thus far.
We trained all models using 10 , 000 examples per task, comparing to elastic weight consolidation (EWC; Kirkpatrick et al., 2017) and regular gradient descent training.
In all cases we rely on a two layer MLP and use Adam (Kingma & Ba, 2014) as the optimiser.
The EWC penalty cost was chosen using a grid search, as was the local MbPA learning rate
Figure 3 compares our approach with that of the baselines.
For this particular task we worked directly on pixels as our embedding, i.e. f γ is the identity function, and explored regimes where the episodic memory is small.
A key takeaway of this experiment is that once a task is catastrophically forgotten, we ﬁnd that only a few gradient steps on carefully selected data from memory are sufﬁcient to recover performance, as MbPA does.
Considering the number of updates required to reach the solution from random initialisation, this fact itself might seem surprising.
MbPA provides a principled and effective way of performing these updates.
The naive approach of performing updates on memories chosen at random from the entire memory is considerably less useful.
We outperformed the MLP, and were superior to EWC for all but one memory size (when storing only a 100 examples per task).
Further, the performance of our model grew with the number of examples stored, ceteris paribus.
Crucially, our memory requirements are much lower than that of
EWC, which requires storing model parameters and Fisher matrices for all tasks seen so far.
Unlike
EWC we do not store any tasks identiﬁers, merely appending the memory with a few examples.
Further, MbPA does not use knowledge of exact task boundaries or identities of tasks switched to, unlike EWC and other methods.
This allows for frequent switches that would otherwise hamper the
Fisher calculations needed for models like EWC.
Our method can be combined with any other algorithm such as standard replay from the memory buffer or EWC, providing further improvement.
In Figure 3 (right) we combine MbPA and EWC.
4. I NCREMENTAL LEARNING : S HIFTS IN T ASK L ABEL D ISTRIBUTIONS
The goal of this section was to evaluate the model in the context of incremental learning.
We considered a classiﬁcation scenario where a model pre-trained on a subset of classes, was introduced to novel, previously unseen classes.
The aim was to incorporate the new related knowledge, as quickly as possible, while preserving knowledge from the previous set.
This was as opposed to the continual learning problem where there are distinct tasks without the ability to revisit old data.
[Figure]
Speciﬁcally we considered the problem of image classiﬁcation on the ImageNet dataset (Russakovsky et al., 2015).
As a parametric model we used a ResnetV1 model (He et al., 2016).
This was pre-trained on a random subset of the ImageNet dataset containing half of the classes.
We then presented all 1000 classes and evaluated how quickly the network can acquire this knowledge (i.e. perform well across all 1000 classes).
For MbPA, we used the penultimate layer of the network as the embedding network f γ , forming the key h and query q for our episodic memory M .
The last fully connected layer was used to initialise the parametric model g θ .
MbPA was applied at test time, using RMSprop with a local learning rate α M and the number of optimisation steps (as in Algorithm 1) tuned as hyper-parameters.
A natural baseline was to simply ﬁne-tune the last layer of the parametric model with the new training set.
We also evaluated a mixture model, combining the classiﬁcations of the parametric model and the non-parametric model at decision level in the following manner:
[Equation]
where the parameter λ ∈ [0 , 1] controls the contribution of each model (this model was proposed by
Grave et al.
in the context of language modelling).
We created ﬁve random splits in new and old classes.
Hyperparameters were tuned for all models using the ﬁrst split and the validation set, and we report the average performance on the remaining splits evaluated on the test set.
Figure 4 shows the test set performance for all models, split by new and old classes.
While the mixture model provides a large improvement over the plain parametric model, MbPA signiﬁcantly outperforms both of them both in speed and performance.
This is particularly noticeable in the new classes, where MbPA acquires knowledge from very few examples.
Table 1 shows a quantitative analysis of these observations.
After around 30 epoches the parametric model matches the performance of MbPA.
In the appendix we explore sensitivity of MbPA on this task to various hyperparameters (memory size, learning rate).
4.2. U NBALANCED DATASETS
We further explored the incremental introduction of new classes, speciﬁcally in the context of unbalanced datasets.
Most real world data are unbalanced, whereas standard datasets (like ImageNet) are artiﬁcially balanced to play well with deep learning methods.
We replicated the setting from the ImageNet experiments in the previous section, where new classes were introduced to a pre-trained model.
However, we only showed a tenth of the data for half the new classes and all data for the other half.
We report performance on the full balanced validation set.
Once again, we compared the parametric model with MbPA and a memory based mixture model.
Results are summarised in Figure 5 (left).
After 20 epochs of training, MbPA outperformed both baselines, with a wider gap in performance than the previous experiment.
Further, the mixture model, though equipped with memory, did signiﬁcantly worse than MbPA, leading us to conclude that the inductive bias in the local adaptation process was well suited to deal with data scarcity.
[Table]
[Figure]
4. L ANGUAGE MODELLING : D OMAIN S HIFTS
Finally we considered how MbPA can be used at test time to further improve the performance of language modelling.
Given the general formulation of MbPA, this could be applied to any problem where there is a shift in distribution at test time — we focus on language modelling, where using recent information has proved promising, such as neural cache and dynamic evaluation (Grave et al., 2016; Krause et al., 2017).
We considered two datasets with established performance benchmarks, Penn Treebank (PTB; Marcus et al., 1993) and WikiText-2 (Merity et al., 2016).
We pre-trained an LSTM and apply MbPA to the weights and biases of the output softmax layer.
The memory stores the past LSTM outputs and associated class labels observed during evaluation.
Full model details and hyper-parameters are detailed in Appendix B.
Penn Treebank is a small text corpus containing 887 , 521 train tokens, 70 , 390 validation tokens, and 78 , 669 test tokens; with a vocabulary size of 10 , 000 .
The LSTM obtained a test perplexity of 59 .and
this dropped by 4 . 3 points when interpolated with the neural cache.
When we interpolated an
LSTM with MbPA we were able to improve on the LSTM baseline by 5 .
3 points (an additional one from the cache model).
We also attempted a dynamic evaluation scheme in a similar style to
Krause et al.
(2017), where we loaded the Adam optimisation parameters obtained during training and evaluated with training of the LSTM enabled, using a BPTT window of 5 steps.
However we did not manage to obtain gains above 1 perplexity from baseline, and so we did not try it for WikiText-2.
WikiText-2 is a larger text corpus than PTB, derived from Wikipedia articles.
It contains 2 , 088 ,train tokens, 217 , 646 validation tokens, and 245 , 569 test tokens, with a vocabulary of 33 , 278 .
Our
LSTM baseline obtained a test perplexity of 65 . 9 , and this is improved by 14 . 6 points when mixed with a neural cache.
Combining the baseline LSTM with an LSTM ﬁt with MbPA we see a drop of 9 . 9 points, however the combination of all three models (LSTM baseline + MbPA + cache) produced
[Figure]
the largest drop of 15 . 9 points.
Comparing the perplexity word-by-word between LSTM + cache and LSTM + cache + MbPA, we see that MbPA improves predictions for rarer words (Figure 8).
PTB
WikiTextValid
Test ∆Test
Valid
Test ∆Test
CharCNN (Zhang et al., 2015) 78.
Variational LSTM (Aharoni et al., 2017) 61.
LSTM + cache (Grave et al., 2016) 74.
LSTM (Melis et al., 2017) 60.
AWD-LSTM (Merity et al., 2017) 60.
AWD-LSTM + cache (Merity et al., 2017) 53.
AWD-LSTM
[Table]
C ONCLUSION
We have described Memory-based Parameter Adaptation (MbPA), a scheme for using an episodic memory structure to locally adapt the parameters of a neural network based upon the current input context.
MbPA works well on a wide range of supervised learning tasks in several incremental, lifelong learning settings: image classiﬁcation, language modelling.
Our experiments show that MbPA improves performance in continual learning experiments, comparable to or in many cases exceeding the performance of EWC.
We also demonstrated that MbPA allows neural networks to rapidly adapt to previously unseen classes in large-scale image classiﬁcation problems using the ImageNet dataset.
Furthermore
[Table]
MbPA can use the local, contextual updates from memory to counter and alleviate the effect of imbalanced classiﬁcation data
[Table]
where some new classes are over-represented at train time whilst others are underrepresented.
Finally we demonstrated on two language modelling tasks that
MbPA is able to adapts to shifts in word distribution common in language modelling tasks, achieving signiﬁcant improvements in performance compared to LSTMs and building on methods like neural cache (Grave et al., 2016).
A CKNOWLEDGMENTS
We would like to thank Gabor Melis for providing the LSTM baselines on the language tasks.
We would also like to thank Dharshan Kumaran, Jonathan Hunt, Olivier Tieleman, Koray Kavukcuoglu,
Daan Wierstra, Sam Ritter, Jane Wang, Alistair Muldal, Nando de Frietas, Tim Harley, Jacob Menick and Steven Hansen for many helpful comments and invigorating discussions
A M B PA H YPERPARAMETERS FOR I NCREMENTAL L EARNING I MAGE N ET T ASK
MbPA was robust and the inductive bias of the MbPA correction adapts the performance of the model on novel classes.
This is shown in Figure 6 (right) where MbPA manages to achieve high performance almost at the same rate, regardless of the learning rate of the underlying parametric component.
In Figure 6 (left) we explore the inﬂuence in performance when changing the size of the episodic memory.
We can see that the performance on the new classes is more sensitive to this parameter but it quickly saturates after about 400,000 entries.
We repeat the above experiment by changing now the number of neighbours retrieved.
The results are shown in Figure 7.
We can observe that using more neighbours is better, but again, performance saturates quickly after 50 neighbours
[Figure]
[Figure]
B M ODEL D ETAILS L ANGUAGE M ODELLING T ASKS
For both datasets we used a single-layer LSTM baseline trained with Adam (Kingma & Ba, 2014) using the regularisation techniques described in Melis et al.
In this application of MbPA the test set is small (e.g. < 80 , 000 words for PTB), and so it was easy to overﬁt to the retrieved points.
To remedy this, we tuned an L2 penalty
β
|| θ x − θ || 2 term in our MbPA loss (7), where θ were the parameters derived from the training set and β was a scalar hyper-parameter.
We swept over the following hyper-parameters:
• Memory size: N ∈ { 500 , 1000 , 5000 }
• Nearest neighbours: K ∈ { 256 , 512 } • Cache interpolation: λ cache ∈ { 0 , 0 .
05 , 0 . 1 , 0 . 15 } • MbPA interpolation: λ mbpa ∈ { 0 , 0 . 05 , 0 . 1 , 0 . 15 } • Number of MbPA optimisation steps: T ∈ { 1 , 5 , 10 } • MbPA optimization learning rate: α ∈ { 0 . 01 , 0 . 1 , 0 . 15 , 0 . 2 , 0 .
Where memory size refers to both the MbPA memory size, and the size of the neural cache for comparison, and the λ interpolation parameters refer to the mixing of model outputs, alike to Eq.
The optimal parameters were: N = 5000 , K = 256 , λ cache = 0 .
15 , λ mbpa = 0 .
1 , T = 1 , α = 0 .
For Penn Treebank, we used a pre-trained LSTM baseline containing roughly 10 M parameters with a hidden size of 1194 and a word embedding size of 268 .
For WikiText-2, we used a pretrained LSTM baseline containing roughly 24 M parameters with a hidden size of 1 , 853 and a word embedding size of 241 .
C C OMPARISON OF C ACHE VS M B PA FOR W IKI T EXT -
The comparative beneﬁt of MbPA is investigated, when combined with the LSTM + cache model.
By computing the perplexity on a per-word basis and comparing whether the inclusion of MbPA improves (lowers) the perplexity, we can understand what types of words are better predicted.
Anecdotal samples were not sufﬁcient to understand the trend, however when the words were bucketed by their training frequency, we see a tend of improved performance for less frequent words.
This improved performance for rare words may be because the cache model has a prior to boost all recent words.
Speciﬁcally, the cache probabilities are obtained from summing the attention for each instance of a word in memory, and so frequently occurring recent words that are not very contextually relevant will still be boosted.
As MbPA does not do this, it appears to be more sensitive to infrequently occurring words.
[Figure]
D MAP INTERPRETATION OF MBPA AND DERIVATION OF CONTEXTUAL UPDATE
Let x c correspond to the input of the h c , v c key-value pair in the context C of a given input x .
In other words, h c was computed by feeding x c to the embedding network.
Then the posterior given this pair and the parameter obtained after training θ , can be written as:
[Equation]
If we maximise the posterior over the context C with respect to θ x .
[Equation]
Let log p ( θ x | θ ) ∝ − ||
θ x − θ ||
22 2 α M (i.e. a Gaussian prior on θ x centred at θ ) be thought as a regularisation term that prevents the local adaptation to move θ x too far from θ , preventing overﬁtting.
Another interpretation of (7) is that when the prior is taken to be a Gaussian, it is a form of elastic weight regularisation (similar to Kirkpatrick et al.
and the second term corresponds to the log likelihood of θ x on the data in the context C .
This can also be seen as posterior sharpening (Fortunato et al., 2017), where we can think of the second term as an approximation of log p ( y t | x t , θ ) .
Thus a view of MbPA is it is a form of local elastic weight consolidation on a context dataset C .
Equation (7) does not have a closed form solution, and requires ﬁtting a large number of parameters at inference time.
This can be costly and susceptible to overﬁtting.
We can avoid this problem by simply adapting the reference parameters θ .
Speciﬁcally, we perform a ﬁxed number of gradient descent steps (or any of its popular variants) to minimise (7).
One step of gradient descent to the loss in (7) with respect to θ x yields
[Equation]
where β is a scalar hyper-parameter.
These adapted parameters are used for output computation but discarded thereafter.
E A TTENTION AS A SPECIAL CASE OF M B PA
Let C = { ( w i , h i , v i ) } ki
=1 be the neighbourhood retrieved from memory given a query q .
The likelihood prediction based on attention is given by
[Equation]
where the Kronecker δ is one when the equality holds and zero otherwise.
We now show how the attention-based prediction given in (9) can be seen as particular case of local adaptation.
For classiﬁcation with c classes, we parameterise p mem via its logits, z ∈ R c , with p mem ( v | q ) =
softmax ( z ) .
One good candidate z is the one that is the most consistent with context C .
Speciﬁcally, the logit vector that minimises the weighted average negative log-likelihood (NLL) of the memories in context C :
[Equation] [Equation]
The attention weights scale the importance of each memory in the neighbour given its similarity to the query.
This matches the loss (7) (ignoring the prior term).
If we differentiate the above equation with respect to a z j and set to zero, we obtain exactly the same expression as in (9).
Effectively, in (10) we are ﬁtting a constant function to the context retrieved from the episodic memory.
This is a particular case of a local likelihood model (Loader, 2006).
The update also is the same as applying a k-nn, see Figure 2.
Note that this interpretation is not limited to classiﬁcation tasks, the exact same reasoning (and result) could be done for a regression task, simply by changing the loss function to be Mean Squared
Error (MSE).
In this context, we can think of MbPA as a generalisation of the attention mechanism, in which the function used for the local ﬁtting is given by the output network.
Moreover, the parameters of the model are used as a prior for solving the local ﬁtting problem and only change slightly to prevent overﬁtting.