MEMO: A Deep Network for Flexible Combination of Episodic Memories

1 INTRODUCTION . During our every day life we need to make several judgments that require connecting facts which were not experienced together , but acquired across experiences at different points in time . For instance , imagine walking your daughter to a coding summer camp and encountering another little girl with a woman . You might conclude that the woman is the mother of the little girl . Few weeks later , you are at a coffee shop near your house and you see the same little girl , this time with a man . Based on these two separated episodes you might infer that there is a relationship between the woman and the man . This flexible recombination of single experiences in novel ways to infer unobserved relationships is called inferential reasoning and is supported by the hippocampus ( Zeithamova et al. , 2012 ) . Interestingly , it has been shown that the hippocampus is storing memories independently of each other through a process called pattern separation ( Yassa & Stark , 2011 ; Marr et al. , 1991 ) . The reason hippocampal memories are kept separated is to minimize interference between experiences , which allows us to recall specific events in the form of ’ episodic ’ memories ( Eichenbaum & Cohen , 2004 ; Squire et al. , 2004 ) . Clearly , this separation is in conflict with the above mentioned role of the hippocampus in generalisation – i.e . how can separated memories be chained together ? Interestingly , a recent line of research ( Kumaran & McClelland , 2012 ; Banino et al. , 2016 ; Schapiro et al. , 2017 ; Koster et al. , 2018 ) sheds lights on this tension by showing that the integration of separated experiences emerges at the point of retrieval through a recurrent mechanism . This allows ∗1DeepMind , London , UK . 2 CoMPLEX , University College London , London , UK . 3 Department of Cell and Developmental Biology , University College London , London , UK . Correspondece should be sent to abanino @ google.com , adriap @ google.com , cblundell @ google.com † , ‡ These authors contributed equally . multiple pattern separated codes to interact , and therefore support inference . In this paper we rely on these findings to investigate how we can take inspiration from neuroscience models to investigate and enhance inferential reasoning in neural networks . Neural networks augmented with external memory , like the Differential Neural Computer ( Graves et al. , 2016 , DNC ) , and end to end memory networks ( Sukhbaatar et al. , 2015 , EMN ) have shown remarkable abilities to tackle difficult computational and reasoning tasks . Also , more powerful attention mechanisms ( Vaswani et al. , 2017 ; Dehghani et al. , 2018 ) or the use of context ( Seo et al. , 2016 ) have recently allowed traditional neural networks to tackle the same set of tasks . However , some of these tasks – e.g . bAbI ( Weston et al. , 2015 ) – present repetitions and commonalities between the train and the test set that neural networks can exploit to come up with degenerate solutions . To overcome this limitation we introduced a new task , called Paired Associative Inference ( PAI - see below ) , which is derived from the neuroscientific literature ( Bunsey & Eichenbaum , 1996 ; Banino et al. , 2016 ) . This task is meant to capture the essence of inferential reasoning – i.e . the appreciation of distant relationships among elements distributed across multiple facts or memories . PAI is fully procedurally generated and so it is designed to force neural networks to learn abstractions to solve previously unseen associations . We then use the PAI task , followed by a task involving finding the shortest path and finally bAbi to investigate what kind of memory representations effectively support memory based reasoning . The EMN and other similar models ( Sukhbaatar et al. , 2015 ; Santoro et al. , 2017 ; Pavez et al. , 2018 ) have used fixed memory representations based on combining word embeddings with a positional encoding transformation . A similar approach has been recently implemented by current state of the art language model ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ) . By contrast our approach , called MEMO , retains the full set of facts into memory , and then learns a linear projection paired with a powerful recurrent attention mechanism that enable greater flexibility in the use of these memories . MEMO is based on the same basic structure of the external memory presented in EMN ( Sukhbaatar et al. , 2015 ) , but its new architectural components can potentially allow for flexible weighting of individual elements in memory and so supporting the form of the inferential reasoning outlined above . Next , we tackle the problem of prohibitive computation time . In standard neural networks , the computation grows as a function of the size of the input , instead of the complexity of the problem being learnt . Sometimes the input is padded with a fixed number of extra values to provide greater computation ( Graves et al. , 2016 ) , in other cases , input values are systematically dropped to reduce the amount of computation ( e.g. , frame dropping in reinforcement learning ( Mnih et al. , 2016 ) ) . Critically , these values are normally hand tuned by the experimenter ; instead , here we are interested in adapting the amount of compute time to the complexity of the task . To do so we drawn inspiration from a model of human associative memory called REMERGE ( Kumaran & McClelland , 2012 ) . In this model , the content retrieved from memory is recirculated back as the new query , then the difference between the content retrieved at different time steps in the re-circulation process is used to calculate if the network has settled into a fixed point , and if so this process terminates . To implement this principle in a neural network , we were inspired by techniques such as adaptive computation time ( Graves , 2016 ) . In our architecture , the network outputs an action ( in the reinforcement learning sense ) that indicates whether it wishes to continue computing and querying its memory , or whether it is able to answer the given task . We call this the halting policy as the network learns the termination criteria of a fixed point operator . Like ACT , the network outputs a probability of halting , but unlike ACT , the binary halting random variable is trained using REINFORCE ( Williams , 1992 ) . Thus we use reinforcement learning to adjust weights based upon the counterfactual problem : what would be the optimal number of steps of computation , given a particular number of steps was taken this time ? The use of REINFORCE to perform variable amount of computation has been investigated already ( e.g . Shen et al. , 2017 ; Louizos et al. , 2017 ) however our approach differs in that we added an extra term to the REINFORCE loss that , by exploiting the mathematical properties of binary random variables , naturally minimizes the expected number of computation steps . Thus we directly encourage our network to explicitly prefer representations and computation that minimize the amount of required computation . To sum up , our contributions are : 1 . A new task that stresses the essence of reasoning — i.e . the appreciation of distant relationships among elements distributed across multiple facts . 2 . An in depth investigation of the memory representation that support inferential reasoning , and extensions to existing memory architectures that show promising results on these reasoning tasks . 3 . A REINFORCE loss component that learn the optimal number of iterations required to learn to solve a task . 4 . Significant empirical results on three tasks demonstrating the effectiveness of the above two contributions : paired associative inference , shortest path finding , and bAbI ( Weston et al. , 2015 ) . 2 METHODS . 2.1 RECAPITULATING END-TO-END MEMORY NETWORKS . We begin by describing End-to-End Memory Networks ( Sukhbaatar et al. , 2015 , EMN ) , as a reminder of this architecture , to introduce notation and nomenclature , and also as a contrast to our work . We will focus on the multilayer , tied weight variant of EMN as this most closely resembles our architecture . The set up used for the rest of the paper is as follows : given a set of knowledge inputs { xi } Ii=1 = { { xi1 , xi2 , ... , xiS } } Ii=1 , and a query or question q = { q1 , q2 , ... , qS } ∈ RS , the network must predict the answer a. I represents the length of the knowledge input sequence , and S is the length of each input sentence ; for instance , in bAbI ( Weston et al. , 2015 ) , I will be the number of stories and S is the number of words in each sentence in a story . xis will be the word in position s in the sentence , in the ith story and will be a O-dimensional one hot vector encoding one of O possible input words . EMN embeds each word and sums the resulting vectors : ki = ∑ s ls ·Wkxis ( 1 ) vi = ∑ s Wvxis ( 2 ) q0 =Wqq ( 3 ) where Wk , Wv ∈ Rd×O , Wq ∈ Rd×S are embedding matrices for the key , values and query , respectively . Also ls is a positional encoding column vector ( as defined in ( Sukhbaatar et al. , 2015 ) ) , · represent an element wise multiplication and O is the size of the vocabulary . At each step t , EMN calculates the vector of weights over the memory elements ki and produces the output . Let K be the I × d matrix formed by taking ki as its rows , and similarly V formed by vi as rows , then : wt = softmax ( Kqt ) ( 4 ) qt+1 = wtV +Wqvqt ( 5 ) at = softmax ( Waqt+1 ) ( 6 ) where wt ∈ RI are weights over the memory slots , Wqv , Wa ∈ Rd×d is a linear mapping relating the query at the previous step to the current one , qt+1 is the query to be used at the next step , and at is the answer ( usually only produced right at the end ) . EMN is trained via a cross entropy loss on at at the final step . 2.2 MEMO . MEMO embeds the input differently . First , a common embedding ci , of size S × dc , is derived for each input matrix xi ∈ RS×O : ci = xiWc ( 7 ) where Wc ∈ RO×dc . Then each of these embeddings is adapted to either be a key or a value . However , contrary to EMN , we do not use hand-coded positional embeddings , instead the words in each sentence and their one-hot encoding in xi , embedded as ci , are combined and then this vector is passed trough a linear projection followed by an attention mechanism ( explained in detail below ) . This allows ci to capture flexibly any part of the input sentence in xi . MEMO uses multiple heads to attend to the memory following ( Vaswani et al. , 2017 ) . Each head has a different view of the same common inputs ci . Let H denote the total number of heads , and h index the particular head , then for each h ∈ { 1 , . . . , H } we have : k ( h ) i =W ( h ) k vec ( ci ) ( 8 ) v ( h ) i =W ( h ) v vec ( ci ) ( 9 ) q ( h ) 0 =W ( h ) q q ( 10 ) where W ( h ) k , W ( h ) v ∈ Rd×Sdc and W ( h ) q ∈ Rd×S are embedding matrices for the key , values and query respectively . vec ( c ) means flattening the matrix c into a vector with the same number of elements as the original matrix , and vec−1 ( v ) is the reverse operation of a vector v into a matrix such that vec−1 ( vec ( c ) ) = c. The result is three d-dimensional vectors k ( h ) i , v ( h ) i and q ( h ) 0 . Keeping each item separated into memory allow us to learn how to weight each of these items when we perform a memory lookup . This contrasts with the hand-coded positional embeddings used in EMN ( Sukhbaatar et al. , 2015 ) and updated recently in Vaswani et al . ( 2017 ) , and proved critical for enabling the flexible recombination of the stored items . The attention mechanism used by MEMO also differs from that shown above for EMN . Firstly , it is adapted to use multi-head attention . Secondly , we use DropOut ( Srivastava et al. , 2014 ) and LayerNorm ( Ba et al. , 2016 ) to improve generalisation and learning dynamics . Let K ( h ) ∈ RI×d denote the matrix formed by taking each k ( h ) i as a row , and V ( h ) be the matrix formed by taking each v ( h ) i as a row . In contrast , let Qt ∈ RH×d be the matrix formed by taking each q ( h ) t as the rows . The attention mechanism then becomes : h ( h ) t = 1√ d WhK ( h ) q ( h ) t ( 11 ) w ( h ) t = DropOut ( softmax ( h ( h ) t ) ) ( 12 ) q ( h ) t+1 = w ( h ) t V ( h ) ( 13 ) Qt+1 = LayerNorm ( vec−1 ( Wqvec ( Qt+1 ) ) +Qt ) ( 14 ) at = softmax ( WaDropOut ( relu ( Wqavec ( Qt+1 ) ) ) ) ( 15 ) where Wh ∈ RI×I , Wq ∈ RHd×Hd are matrices for transforming the logits and queries respectively . Wa ∈ RO×da and Wqa ∈ Rda×Hd are the matrices of the output MLP that produces the answer at . It is worth noting , that even though our attention mechanism uses some of the feature implemented in Vaswani et al . ( 2017 ) – i.e . normalization factor sqrt ( d ) and multiheading , it differs from it because rather than doing self-attention it preserves the query separated from the keys and the values . This aspect is particularly important in terms of computational complexity , in that MEMO is linear with respect to the number of sentences of the input , whereas methods relying on self-attention ( e.g . Dehghani et al. , 2018 ) have quadratic complexity ( see Appendix E ) .

Summary: This paper presents a new task (paired associate inference), drawn from cognitive psychology, which requires linking many pieces of information together to make inferences with long range dependencies. Experimental results show that standard memory architectures fail on these tasks. To redress this, the paper proposes a new memory architecture with several new features that allow for much better performance on the paired associate task. Finally, the paper undertakes systematic experiments on more traditional domains like shortest path problems, showing that the new architecture achieves modest improvements.