Learning What to Remember: Long-term Episodic Memory Networks for Learning from Streaming Data

1 INTRODUCTION . In a real world machine learning problem , it is common to process a tremendous amount of sequential data ( e.g . dialogues , videos , and medical records ) . To model such sequential data , researchers often use recurrent neural networks ( RNNs ) because of its capability to process sequentially dependent data . Recently , popular variants of RNNs such as Long Short-Term Memory ( LSTM , Hochreiter & Schmidhuber ( 1997 ) ) and Gated Recurrent Unit ( GRU , Chung et al . ( 2014 ) ) have achieved impressive performances on a number of tasks , such as machine translation , image annotation , and question answering . Nevertheless , despite such impressive achievements of the recurrent neural networks , it is still challenging to model long sequential data such as dialogues since RNNs have a limited memory capacity . The only memory that LSTM or RNN contains is a fixed dimensional vector . The LSTM or RNN performs memorization by adding to , and erasing , this single vector . This limited memory capacity can be problematic when dealing with a long sequence . To overcome this scalability problem , researchers have proposed neural network architectures that contain external memory module , including Neural Turing machine ( NTM , Graves et al . ( 2014 ) ) and End-to-End Memory Networks ( Sukhbaatar et al . ( 2015 ) ) . These memory-augmented networks have mechanisms to write representation to the memory cells and read from them . The addressing of the memory cells can be done either by an index to perform sequential read/write or by contents to perform a content-based access . The memory-augmented networks have been shown to be useful for memory-oriented tasks such as copying , pathfinding , and question answering . Yet , they are mostly targeting toy problems , using the external memory as a working memory in the human memory mechanism , and need to manage a relatively small number of memory entries . Thus , there has been less focus on how to efficiently manage the memory . However , when the network needs to process data that are too large to fit into the memory , we need to carefully manage the memory cell such that the external memory module contains the most informative pieces of data . In other words , we should be able to determine which memory cell is unlikely to be used in the future so it can be replaced with the incoming data when memory is full . To this end , we consider this memory management problem as a learning problem and train an agent by reinforcement learning , such that it learns to keep the most important information in memory and evict unimportant ones . How can we then select which memory entries are most important ? To achieve this goal , we propose to train the memory module itself using reinforcement learning to delete the most uninformative memory entry in order to maximize its reward on a future task . However , this is a challenging task since for most of the time , this scheduling should be performed without knowing which task will arrive when . Thus , deciding which memory to keep and which to erase should be done by considering relative generic importance of the memory entries . To tackle this challenge , we implemented the memory retention mechanism using a spatio-temporal recurrent neural network , that can learn relative importance among the memory entries as well as their historic importance . We refer to this memory-augmented networks with spatio-temporal retention mechanism as Long Term Episodic Memory Networks ( LEMN ) . LEMN can perform selective memorization to keep a compact set of important pieces of data that will be useful for future tasks in lifelong learning scenarios . We validate our proposed retention mechanism against naive scheduling method as well as RL-based scheduling on three different tasks , namely path-finding , episodic question answering and long question answering , against which it significantly outperforms . Our contribution is twofold : • We consider a novel task of learning from streaming data , where the size of the memory is significantly smaller than the length of the data stream . • To implement a retention mechanism , we propose a retention agent that can be integrated with existing external memory neural networks , which is trained with reinforcement learning to keep memory cells of general importance . 2 RELATED WORK . Memory-augmented networks Graves et al . ( 2014 ) propose Neural Turing Machine ( NTM ) , which has a memory and a controller that reads from and writes to memory . The controller composed of read heads and write heads uses soft-attention mechanism to access memory for differentiable memory access . The NTM has two addressing mechanisms : content-based addressing and locationbased addressing . Location-based addressing allows a data sequence to be stored in consecutive memory cells to preserve its sequential order . However , once the write head accesses a distant memory cells , the sequential order of information in consecutive memory cells is no longer preserved . Graves et al . ( 2016 ) propose Differentiable Neural Computer that extends the NTM to address the issue by introducing a temporal link matrix . As it is costly to write into memory while preserving its sequential order , memory-augmented networks with write mechanism are mostly used in cases where it is not necessary to track which sequential order the memory has been written ( Santoro et al. , 2016 ; Vinyals et al. , 2016 ; Kaiser et al. , 2017 ; Kim et al. , 2018 ) . Unlike NTM , End-to-End Memory Networks ( MemN2N ) ( Sukhbaatar et al. , 2015 ) and Dynamic Memory Networks ( DMN+ ) ( Xiong et al. , 2016 ) do not have write mechanism but store all the inputs into memory . For this reason , the sequential order of the inputs are preserved at no extra cost ; thus they are more suitable for episodic question answering problems such as bAbI tasks ( Weston et al. , 2015 ) . In addition , they have sophisticated read mechanisms that allow to reason through multiple hops ( or passes in DMN+ ) . Seo et al . ( 2016 ) propose more advanced soft-attention based read mechanism , Bi-Directional Attention Flow ( BiDAF ) , which obtains impressive performance on a difficult reading comprehension dataset , Stanford Question Answering Dataset ( SQuAD , Rajpurkar et al . ( 2016 ) ) . Oh et al . ( 2016 ) extend MemN2N to train deep reinforcement learning agents , Memory Q-Network ( MQN ) and Feedback Recurrent Memory Q-Network ( FRMQN ) , to solve 3D Mazes . However , such neural networks have the limitation that the size of external memory should be large enough to store all the data . Memory Retention Policy Most of the existing approaches ( MemN2N , DMN+ , BiDAF and FRMQN ) do not consider the case where memory becomes full , and simply truncate the sequence of data to the size of memory if it is longer than the memory size . Many of mutable external memory neural networks employ least-recently-used ( LRU ) based memory retention policy , which overwrites new data into the least used memory cell . This policy , although more reasonable than FIFO , is still limited as it is a hand-crafted rule and does not consider actual long-term dependencies between a memory entry and the task . This is a critical limitation in lifelong learning setting , since some of the memory entries written in the long past and have not been accessed for long may still be necessary to respond to queries that arrive in the far future . Our work , on the other hand , is able to learn such long-term dependencies . DNTM ( Gülçehre et al. , 2016 ) learns where to overwrite using reinforcement learning as done in our work . Yet , it only considers the pairwise relationships between the current data instance and each individual memory entries , while our model learns both relative importance and historic importance of a memory entry using a spatio-temporal RNN architecture . 3 LEARNING WHAT TO REMEMBER FROM STREAMING DATA . We consider the problem of learning from a long data stream that contains large portion of unimportant , noisy data ( e.g . routine greetings in dialogs ) with limited memory . Formally , an agent A takes as input a streaming data ( e.g . sentences or images ) X = { x1 , · · · , xT } and manages an external memory M = [ m1 , · · · , mN ] , mi ∈ Rd where T N . Thus , the agent should decide which memory cell to evict to store incoming data . To this end , the agent should learn the relative importance of memory entries for a future task . In traditional reinforcement learning scheme , we can formulate this problem as learning the policy π ( at|st ) to maximize a return R , where action at is a memory cell to erase , and state st = [ Mt ; xt ] , where Mt is the current memory and xt is the input at time t. The agent appends xt to the memory Mt until it reaches the maximum size N . If the memory is full , the agent erases one memory cell based on the policy π ( mi|Mt , xt ) to append xt . Thus , we can consider 1 − π ( mi|Mt , xt ) as the retention value of each memory . At arbitrary time step tf , it encounters a task T ( e.q . question answering ) with a reward RT , whose reward is defined by the specific task . For QA task , the reward will be 1 if it generated a correct answer and−1 otherwise . The instance xt which arrive at timestep t can be in any data format , and is transformed into a memory vector representation et ∈ Rd to be stored in memory . 3.1 MEMORY-RETENTION MECHANISMS . In this section , we describe three different memory-retention mechanisms in detail . We first encode input xt and each memory cell mt , i to a vector representation as follows : ct = ψ ( xt ) et , i = ϕ ( mt , i ) where ct , et , i ∈ Rd , ψ ( · ) can be an input embedding similar to RNN controller in ( Gülçehre et al. , 2016 ) or position encoding in ( Sukhbaatar et al. , 2015 ) and ϕ ( · ) can be a memory encoding layer of the base external memory neural network . For example , MemN2N ( Sukhbaatar et al. , 2015 ) uses a bag of words and embedding matrices to convert to memory representation . Input-Matching LEMN This is the simplest RL-based retention mechanism that is similar to DNTM from Gülçehre et al . ( 2016 ) and utilizes the memory usage information , which we refer to as Input-Matching LEMN ( IM-LEMN , Figure 1a ) . The usage of each memory representation mi is computed as the learned similarity between it and the current input xt . As in Gülçehre et al . ( 2016 ) , we least recently used ( LRU ) addressing by computing the exponential moving average vt of the logit zt , the LRU factor γt , and the policy as follows : vt , i = 0.1vt−1 , i + 0.9zt , i γt = σ ( Wγct + bγ ) gt , i = zt , i − γtvt−1 , i π ( mi|Mt , xt ) = softmax ( gt , i ) where Wγ ∈ Rd , bγ ∈ R , softmax ( zi ) is exp ( zi ) / ∑ j exp ( zj ) , σ ( z ) is 1/ ( 1 + exp ( z ) ) . This model estimates the policy based on the similarities between input xt and memory cells mt , i . Spatial LEMN A major drawback of IM-LEMN is that the score of each memory depends only on the input xt . In other words , the score is computed independently between memory cell and input but does not consider its relative importance to other data instances in the memory . To overcome this limitation , we propose Spatial LEMN ( S-LEMN , Figure 1b ) which computes the relative importance of memory cells to its neighbors and other memory cells using a bidirectional GRU as follows : −→e t , i = GRUθfw ( et , i , −→e t , i−1 ) ←−e t , i = GRUθbw ( et , i , ←−e t , i+1 ) ft , i = ReLU ( Wf [ −→e t , i , ←−e t , i ] + bf ) where Wf ∈ R2d×d , bf ∈ Rd , GRUθ is a Gated Recurrent Unit parameterized by θ , [ −→e t , i , ←−e t , i ] is a concatenation of features , ReLU is a rectified linear unit . We use multi-layer perceptron ( MLP ) with scalar output to estimate the policy as follows : ht , i =Whft , i + bh gt , i =Wght , i + bg ( 1 ) π ( mi|Mt , xt ) = softmax ( gt , i ) where Wh ∈ Rd×d/4 , bh ∈ Rd/4 , Wg ∈ Rd/4 , bg ∈ R. Thus , it can compute the general importance of memory cell itself and the relation between its neighbors in contrast to IM-LEMN . Spatio-Temporal LEMN In episodic tasks , the importance of memory also changes over time and tasks . Hence , we propose Spatio-Temporal LEMN ( ST-LEMN , Figure 1c ) that uses a GRU over time to consider the historical importance as well as relative importance of each memory entry . We simply replace the hidden layer in Equation ( 1 ) with a GRU over time as follows : ht , i = GRUθh ( ft , i , ht−1 , i ) gt , i =Wght , i + bg π ( mi|Mt , xt ) = softmax ( gt , i ) Memory Update Using aforementioned memory-retention mechanism , the agent samples the memory cell index i from the probability distribution π ( mi|Mt , xt ) . Then it erases the ith memory cell and appends the input xt as follows : Mt+1 = [ m1 , · · · , mi−1 , mi+1 , · · · , mN , xt ] . No Operation ( NOP ) As done in Gülçehre et al . ( 2016 ) , for IM-LEMN we add a NOP memory cell at the end of memory to consider the case where the input is not written to the memory . For SLEMN and ST-LEMN , we append xt to the place of NOP such that it could be selected for removal . We integrate these retention mechanisms into base memory-augmented neural networks to enable to efficiently maintain its external memory . The details of the base networks are given in the Experiment section . We use Asynchronous Advantage Actor-Critic ( A3C , Mnih et al . ( 2016 ) ) with Generalized Advantage Estimation ( GAE , Schulman et al . ( 2015 ) ) to train all models .

Summary: This paper attempts to study memory-augmented neural networks when the size of the data is too large. The solution is to maintain a fix-sized episodic memory to remember the important data instances and at the same time erase the unimportant instances. To do so, the authors improve the method called DNTM (Gulcehre et al., 2016) by incorporating the similarity between each memory entry besides the similarity between the current data the each memory entry. Experiments show the effectiveness of the proposed method.