Adaptive Posterior Learning: few-shot learning with a surprise-based memory module

1 INTRODUCTION . Consider the following sequential decision problem : at every iteration of an episode we are provided with an image of a digit ( e.g . MNIST ) and an unknown symbol . Our goal is to output a digit Y = X + S where X is the value of the MNIST digit , and S is a numerical value that is randomly assigned to the unknown symbol at the beginning of each episode . After seeing only a single instance of a symbol an intelligent system should not only be able to infer the value S of the symbol but also to correctly generalize the operation associated with the symbol to any other digit in the remaining iterations of that episode . Despite its simplicity , this task emphasizes three cognitive abilities that a generic learning algorithm should display : 1. the algorithm can learn a behaviour and then flexibly apply it to a range of different tasks using only a few context observations at test time ; 2. the algorithm can memorize and quickly recall previous experiences for quick adaptation ; and 3. the algorithm can process these recalled memories in a non-trivial manner to carry out tasks that require reasoning . The first point is commonly described as “ learning to learn ” or meta-learning , and represents a new way of looking at statistical inference ( Schmidhuber , 1987 ; Bengio et al. , 1990 ; Bengio & LeCun , 2007 ) . Traditional neural networks are trained to approximate arbitrary probability distributions with great accuracy by parametric adaptation via gradient descent ( LeCun et al. , 2015 ; Schmidhuber , 2015 ) . After training that probability distribution is fixed and neural networks can only generalize well when the testing distribution matches the training distribution ( Neyshabur et al. , 2017 ) . In contrast , meta-learning systems are trained to learn an algorithm that infers a function directly from the observations it receives at test time . This setup is more flexible than the traditional approach and generalizes better to unseen distributions as it incorporates new information even after the training phase is over . It also allows these models to improve their accuracy as they observe more data , unlike models which learn a fixed distribution . The second requirement - being able to memorize and efficiently recall previous experience - is another active area of research . Storing information in a model proves especially challenging as we move beyond small toy-examples to tasks with higher dimensional data or real-world problems . Current methods often work around this by summarizing past experiences in one lower-dimensional representation ( Hochreiter & Schmidhuber , 1997 ; Kingma & Welling , 2013 ) or using memory modules ( Graves et al. , 2016 ) . While the former approach can produce good results , the representation and therefore the amount of information we can ultimately encode with such models will be of a fixed and thus limited size . Working with neural memory modules , on the other hand , presents its own challenges as learning to store and keep the right experiences is not trivial . In order to successfully carry out the task defined at the beginning of this paper a model should learn to capture information about a flexible and unbounded number of symbols observed in an episode without storing redundant information . Finally , reasoning requires processing recalled experiences in order to apply the information they contain to the current data point being processed . In simple cases such as classification , it is enough to simply recall memories of similar data points and directly infer the current class by combining them using a weighted average or a simple kernel ( Vinyals et al. , 2016 ; Snell et al. , 2017 ) , which limits the models to performing interpolation . In the example mentioned above , more complex reasoning is necessary for human-level generalisation . In this paper we introduce Approximate Posterior Learning ( APL , pronounced like the fruit ) , a self-contained model and training procedure that address these challenges . APL learns to carry out few-shot approximation of new probability distributions and to store only as few context points as possible in order to carry out the current task . In addition it learns how to process recalled experiences to carry out tasks of varying degrees of complexity . This sequential algorithm was inspired by Bayesian posterior updating ( Jaynes , 2003 ) in the sense that the output probability distribution is updated as more data is observed . We demonstrate that APL can deliver accuracy comparable to other state-of-the-art algorithms in standard few-shot classification benchmarks while being more data efficient . We also show it can scale to a significantly larger number of classes while retaining good performance . Finally , we apply APL to the reasoning task introduced as motivation and verify that it can perform the strong generalization we desire . The main contributions of this paper are : • A simple memory controller design which uses a surprise-based signal to write the most predictive items to memory . By not needing to learn what to write , we avoid costly backpropagation through memory which makes the setup easier and faster to train . This design also minimizes how much data is stored , making our method more memory efficient . • An integrated external and working memory architecture which can take advantage of the best of both worlds : scalability and sparse access provided by the working memory ; and all-to-all attention and reasoning provided by a relational reasoning module . • A training setup which steers the system towards learning an algorithm which approximates the posterior without backpropagating through the whole sequence of data in an episode . 2 MODEL . 2.1 ARCHITECTURE . Our proposed model is composed of a number of parts : an encoder that generates a representation for the incoming query data ; an external memory store which contains previously seen representation/ data pairings with writing managed by a memory controller ; and a decoder that ingests the query representation as well as data from the memory store to generate a probability distribution over targets . We describe each of the parts in detail below1 . Encoder The encoder is a function which takes in arbitrary data xt and converts it to a representation et of ( usually ) lower dimensionality . In all our experiments xt is an image , and we therefore choose a convolutional network architecture for the encoder . Architectural details of the encoder used for each of the experiments are provided in the appendix . 1Source code for the model is available at https : //github.com/cogentlabs/apl . Memory store The external memory module is a database containing the stored experiences . Each of the columns corresponds to one of the attributes of the data . In the case of classification , for example , we would store two columns : the embedding em and the true label ym . Each of the rows contains the information for one data point . The memory module is queried by finding the k-nearest neighbors between a query and the data in a given column . The full row data for each of the neighbors is returned for later use . The distance metric used to calculate proximity between the points is an open choice , and here we always use euclidean distance . Since we are not backpropagating through the memory , how do we ensure that the neighbors returned by the querying mechanism contain task-relevant information ? We expect that class-discriminative embeddings produced by the encoder should cluster together in representation space , and therefore should be close in the sense of euclidean distance . While this is not mathematically necessary , in the following sections we will show that APL as proposed does work and retrieve the correct neighbors which means that in practice our intuition holds true . Memory controller We use a simple memory controller which tries to minimize the amount of data points written to memory . Let us define surprise as the quantity associated with the prediction for label yt as S = −ln ( yt ) . Intuitively , this means that the higher the probability our model assigns to the true class , the less surprised it will be . This suggests a way of storing the minimal amount of data points in memory which supports maximal classification accuracy . If a data point is ’ surprising ’ , it should be stored in memory ; otherwise it can be safely discarded as the model can already classify it correctly . How should the memory controller decide whether a point is ’ surprising ’ ? In this work we choose the simplest possible controller : if the surprise is greater than some hyperparameter σ , then that data should be stored in memory . For our experiments , we choose σ ∝ − ln ( N ) where N is the number of classes under classification which means that if the prediction confidence in the correct class is smaller than the probability assigned by a uniform prediction the value should be written to memory . In the appendix we show that after model training model performance is robust to variations in σ , as surprise becomes highly bimodal : a new data point tends to be either highly surprising ( never seen something similar before ) or not very surprising . Conveniently , in the case of classification problems the commonly used cross-entropy loss reduces to our measure of surprise directly , and we therefore use the prediction loss as an input to the memory controller directly . Decoder The decoder takes as input the query representation as well as all the data from the neighbors found in the external memory . We designed a relational feed-forward module with self attention which takes particular advantage of the external memory architecture . In addition we tested two other established decoder architectures : an unrolled relational working memory core and an unrolled LSTM . As all experiments have a classification loss at the end , all the decoders return a vector with logits for the N classes under consideration . Full details of each architecture are provided in the Appendix . • Relational self-attention feed-forward decoder . The relational feed-forward module ( see figure 2 , left ) processes each of the neighbors individually by comparing them with the query , and then does a cross-element comparison with a self-attention module before reducing the activations with an attention vector calculated from neighbor distances . • Relational working memory decoder ( Santoro et al. , 2018 ) The relational working memory module ( figure 2 , right ) takes in the concatenated neighbor embeddings and corresponding label embeddings as its initial memory state . The query is fed a number N times as input to the relational memory core to unroll the computation . • LSTM decoder Finally we also test a vanilla LSTM decoder that takes in the query as the initial memory state and is fed each of the concatenated neighbor embeddings and corresponding label embeddings as its input each time step . 2.2 TRAINING SETUP . Since we are looking for a system which can update its beliefs in an online manner we need a training procedure that reflect this behaviour . We train the system over a sequence of episodes that are composed of sequences of pairs ( xt , yt ) . At the start of every episode the mapping xt → yt is shuffled in a deterministic manner ( the exact details are task dependent and will be outlined in the experiments section ) . The data is then presented to the model sequentially in a random order . The model ’ s memory is empty at the beginning of the episode . At each time step , a batch of examples is shown to the model and a prediction is made . We then measure the instantaneous loss L ( ŷt , yt ) and perform a gradient update step on the network to minimize the loss on that batch alone . The loss is also fed to the memory controller for the network to decide whether to write to memory . In all the experiments below the task is to classify some quantity , therefore we use cross entropy loss throughout . APL learns a sequential update algorithm , that is , it minimizes the expected loss over an episode consisting of a number of data elements presented sequentially . However we don ’ t need to backpropagate through the sequence to learn the algorithm . Rather , the model ’ s parameters are updated to minimize the cross-entropy loss independently at each time step . Therefore the only pressure to learn a sequential algorithm comes from the fact that episodes are kept small so that the decoder is encouraged to read the information coming from the queried neighbors instead of just learning to fit the current episode ’ s label mapping in its weights after a few steps of gradient descent ( which is what happens in the case of MAML ( Finn et al. , 2017 ) ) .

Summary: the authors propose a new algorithm, APL, for a few-shot and a life-long learning based on an external memory module. APL uses a surprise-based signal to determine which data points to store in memory and an attention mechanism to the most relevant points for prediction. The authors evaluate APL on a few-shot classification task on Omniglot dataset and on a number analogy task.