Online Meta-Learning via Learning with Layer-Distributed Memory

1 Introduction . Meta-learning or learning-to-learn is a paradigm that enables models to generalize to a distribution of tasks rather than specialize to just one task [ 1 , 2 ] . When encountering examples from a new task , we would like the model to adapt to the new task after seeing just a few samples . This is commonly achieved via episodic training of deep neural networks , where , in each episode , the network is exposed to a variety of inputs from the same distribution [ 3 , 4 ] , and the distribution shifts over episodes . The ability of deep networks to adapt to a new task within just a few samples or iterations is central to the application of meta-learning methods in few-shot and online learning scenarios [ 5 , 6 ] . A recent surge of interest directed towards meta-learning using neural networks has spurred development of a variety of methods [ 7–9 ] . In a standard episodic training framework , a network must adapt to a sampled task ( or collection of tasks ) and incurs a generalization loss for that task ( or collection ) ; this generalization loss is backpropagated to update the network weights . Methods differ in the underlying architecture and mechanisms they use to support adaptation . Strategies include using gradient descent in an inner loop , storing and updating prototypes , parameterizing update rules by another neural network , and employing neural memory [ 3 , 10–12 ] . Section 2 provides an overview . We focus on memory-based meta-learning , and specifically investigate the organization of neural memory for meta-learning . Motivating this focus is the generality and flexibility of memory-based approaches . Relying on memory for adaptation allows one to cast meta-learning as merely a learning problem using a straightforward loss formulation ( viewing entire episodes as examples ) and standard optimization techniques . The actual burden of adaptation becomes an implicit responsibility of the memory subsystem : the network must learn to use its persistent memory in a manner that facilitates task adaptation . This contrasts with explicit adaptation mechanisms such as stored prototypes . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this implicit adaptation setting , memory architecture plays a crucial role in determining what kind of adaptation can be learned . We experimentally evaluate the effectiveness of alternative neural memory architectures for meta-learning and observe particular advantages to distributing memory throughout a network . More specifically , we view the generic LSTM equations , Wx+Wh−1 , as adaptation induced by hidden states in activation space ( see Figure 1 ) . By distributing LSTM memory cells across the depth of the network , each layer is tasked with generating hidden states that are useful for adaptation . Such a memory organization is compatible with many standard networks , including CNNs , and can be achieved by merely swapping LSTM memory cells in place of existing filters . Our simple approach also contrasts with several existing memory-based meta-learning methods used in both generative and classification tasks [ 13–17 ] . These methods view memory as a means to store and retrieve useful inductive biases for task adaptation , and hence focus on designing better read and write protocols . They typically have a feature extractor that feeds into a memory network that performs adaptation , whereas our architecture makes no such distinction between stages . We test the efficacy of network architectures with distributed memory cells on online few-shot and continual learning tasks as in Santoro et al . [ 13 ] , Ren et al . [ 18 ] and Javed and White [ 6 ] . The online setting is challenging for two reasons : 1 ) It is empirically observed that networks are not well suited for training/adaptation with a batch size of one [ 19 ] ; 2 ) In this setting the model has to adapt to one image at a time step , thus having to deal with a prolonged adaptation phase . For these reasons , we see these tasks as suitable for evaluating the adaptation capabilities of the hidden states generated by the network . We empirically observe that our method outperforms strong gradient-based and prototypical baselines , delineating the efficacy of the local adaptation rule learnt by each layer . Particularly important is the distributed nature of our memory , which allows every network layer to adapt when provided with label information ; in comparison , restricting adaptability to only later network layers delivers far less compelling performance . These results suggest that co-design of memory architecture and metatraining strategies should be a primary consideration in the ongoing development of memory-based meta-learning . We further test our model in a harder online few-shot learning scenario , wherein the corresponding label to a sample arrives after a long delay [ 20 ] . Our method adapts seamlessly , without requiring any changes to the model , while , in this setting , other adaptation strategies struggle . These results highlight promising directions for advancing and simplifying meta-learning by relying upon distributed memory for adaptation . 2 Related Work . Early work on meta-learning introduces many relevant concepts . Schmidhuber [ 21 ] proposes using task specific weights , called fast weights , and weights that are adapted across tasks , called slow weights . Bengio et al . [ 2 ] updates the network via a learning rule which is parameterized by another neural network . Thrun [ 22 ] presents meta-learning in a life-long scenario , where the algorithm accrues information from the past experiences to adapt effectively for the task at hand . Hochreiter et al . [ 23 ] train a memory network to learn its own adaptation rule via just its recurrent states . These high level concepts can be seen in more recent methods . We group current meta-learning methods based on the nature of adaptation strategy and discuss them below . Gradient-based Adaptation Methods . Methods that adapt via gradients constitute a prominent class of meta-learning algorithms [ 9 ] . Model-agnostic meta-learning ( MAML ) [ 4 ] learns an initialization that can efficiently be adapted by gradient descent for a new task . Finn et al . [ 24 ] focus on learning a network that can use experience from previously seen tasks for current task adaptation . They adapt to the current task by using a network that is MAML pre-trained on the samples from the previous task . Nagabandi et al . [ 25 ] , Caccia et al . [ 26 ] perform online adaptation under non-stationary distributions , either by using a mixture model or by spawning a MAML pre-trained network when the input distribution changes . Javed and White [ 6 ] , Beaulieu et al . [ 27 ] employ a bi-level optimization routine similar to MAML , except the outer loop loss is catastrophic forgetting . They thereby learn representations that are robust to forgetting and accelerate future learning under online updates . Memory and Gradient-based Adaptation . Andrychowicz et al . [ 28 ] , Ravi and Larochelle [ 10 ] learn an update rule for network weights by transforming gradients via a LSTM , which outperforms human-designed and fixed SGD update rules . Munkhdalai and Yu [ 29 ] learn a transform that maps gradients to fast ( task specific ) weights , which are stored and retrieved via attention during evaluation . They update slow weights ( across task weights ) at the end of each task . Prototypical Methods . These methods learn an encoder which projects training data to a metric space , and obtain class-wise prototypes via averaging representations within the same class . Following this , test data is mapped to the same metric space , wherein classification is achieved via a simple rule ( e.g. , nearest neighbor prototype based on either euclidean distance or cosine similarity ) [ 5 , 30 , 31 ] . These methods are naturally amenable for online learning as class-wise prototypes can be updated in an online manner as shown by Ren et al . [ 18 ] . Memory-based Adaptation . Santoro et al . [ 13 ] design efficient read and write protocols for a Neural Turning Machine [ 32 ] for the purposes of online few-shot learning . Rae et al . [ 33 ] design sparse read and write operations , thereby making them scalable in both time and space . Ramalho and Garnelo [ 11 ] use logits generated by the model to decide if a certain sample is written to neural memory . Mishra et al . [ 7 ] employ an attention-based mechanism to perform adaptation , and use a CNN to generate features for the attention mechanism . Their model requires storing samples across all time steps explicitly , thereby violating the online learning assumption of being able to access each sample only once . All of these methods mainly focus on designing better memory modules either via using more recent attention mechanisms or by designing better read and write rules to neural memory . These methods typically use a CNN which is not adapted for the current task . Our approach differs from these methods , in that we study efficient organization of memory for both online few-shot learning and meta-learning more generally , and show that as a consequence of our distributed memory organization , the entire network is capable of effective adaptation when provided with relevant feedback . Kirsch and Schmidhuber [ 34 ] introduce an interesting form of weight sharing wherein LSTM cells ( with tied weights ) are distributed throughout the width and depth of the network , however each position has its own hidden state . Further , they have backward connections from the later layers to the earlier layers , enabling the network to implement its own learning algorithm or clone a humandesigned learning algorithm such as backprop . Both our model and theirs implement an adaptation strategy purely using the recurrent states . The difference , however , is in the nature of the adaptation strategy implemented in the recurrent states . Similar to conventional learning algorithms , their backward connections help propagate error from the last layer to the earlier layers . In our architecture , the feedback signal is presented as another input , propagated from the first layer to the last layer . In addition to being used in classification settings , Guez et al . [ 35 ] employ memory-based metalearning approach to perform adaptation for reinforcement learning tasks indicating the generality of using memory as a means for adaptation . Few-shot Semantic Segmentation . Few-shot segmentation methods commonly rely on using prototypes [ 36 , 37 ] , though recent approaches include gradient-based methods analogous to MAML [ 38 ] . The methods that use neural memory typically employ it in final network stages to fuse features of different formats for efficient segmentation : Li et al . [ 39 ] use ConvLSTMs [ 40 ] to fuse features from different stages of the network ; Valipour et al . [ 41 ] to fuse spatio-temporal features while segmenting videos ; Hu et al . [ 42 ] use a ConvLSTM to fuse features of query with the features of support set ; Azad et al . [ 43 ] use a bidirectional ConvLSTM to fuse segmentation derived from multiple scale space representations . We differ from these works in organization , use of , and information provided to memory module : 1 ) Memory is distributed across the network as the sole driver of adaptation ; 2 ) Label information is provided to assist with adaptation . Meta-learning Benchmarks . Caccia et al . [ 26 ] present benchmarks that measure the ability of a model to adapt to a new task , using the inductive biases that it has acquired over solving previously seen tasks . More specifically , the benchmark presents an online non-stationary stream of tasks , and the model ’ s ability to adapt to a new task at each time step is evaluated . Note that they do not measure the model ’ s ability to remember earlier tasks ; they only want the model to adapt well on a newly presented task . Antoniou et al . [ 44 ] present benchmarks for continual few-shot learning . The network is presented a number of few-shot tasks , one after the other , and then is expected to generalize even to the previously seen tasks . This is a challenging and interesting setup , in that , the network has to show robustness to catastrophic forgetting while learning from limited data . However , we are interested in evaluating the online adaptation ability of models , while Antoniou et al . [ 44 ] feed data in a batch setting . We follow experimental setup as in Javed and White [ 6 ] , where in , the model is required to remember inductive biases acquired over a longer time frame when compared to the experimental setup used by Antoniou et al . [ 44 ] .

Summary: The paper proposes to adapt Meta RNNs (memory-based meta-learning) by instantiating multiple layers of LSTMs where some of them are convolutional. The authors refer to this as distributed memory due to each layer having its own LSTM-based memory. They demonstrate good performance in few-shot learning and continual learning.