SEQUENCE MODELLING WITH AUTO-ADDRESSING AND RECURRENT MEMORY INTEGRATING NETWORKS

1 INTRODUCTION . Recurrent neural networks , such as the Long Short-Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) and Gated Recurrent Unit ( GRU ) ( Cho et al. , 2014 ) have shown promising performance for processing sequential data . However , it ’ s known that RNNs suffer from gradient vanishing problem . Moreover , as pointed out by Rae et al . ( 2016 ) , the number of parameters grows proportionally to the square of the size of the hidden units , which carry the historical information . Recent memorybased approaches exhibit potential to address these issues , by decoupling memory capacity from model parameters , and backpropagating the gradients through the memory . Neural Turing Machine ( NTM ) ( Graves et al. , 2014 ) first emerged as a recurrent model that incorporates external memory abilities . NTM maintains a memory matrix , and at every time-step , the network reads and writes ( with erasing ) to the memory matrix using certain soft-attentional mechanism , controlled by an LSTM that produces read and write vectors . NTM and its successor , the Differentiable Neural Computer ( Graves et al. , 2016 ) , have shown success on some algorithmic tasks such as copying , priority sorting and some real-world tasks such as question answering . But one limitation of the NTM is that due to its smooth read and write mechanism , NTM has to do propagations on the entire memory , usually causing huge amount of memory consumption . To this end , Rae et al . ( 2016 ) proposes the Sparse Access Memory ( SAM ) network , by thresholding memory modifications to a sparse subset , i.e . all read and write operations are limited to several memory words . This allows memory and time efficient propagations while maintaining NTM ’ s performance . However , these external memory models have relatively complicated memory addressing mechanisms , making them suitable for only large-scale memory . Moreover , the RNN in these models plays a simple role of being a controller , but the gradient vanishing problem of RNN itself is not given attention . In contrast to these networks , our network uses a light-weight read mechanism to address a small external memory , and relieves the gradient vanishing problem by directly concatenating memory word to output . Inspired by prior memory models , efforts have been made to build a bridge between simple RNNs and complicated memory models . Kurach et al . ( 2015 ) propose the Neural Random-access Machines ( NRAM ) that can manipulate and dereference pointers to an external variable-size random-access memory . Danihelka et al . ( 2016 ) improve LSTM with ideas from Holographic Reduced Representations ( Plate , 2003 ) that enables key-value storage of data . Grave et al . ( 2016 ) propose a method of augmenting LSTM by storing previous ( hidden state , input word ) pairs in memory and using the current hidden state as a query vector to recover historical input words . This method requires no backpropagation through memory and is well-suited to word-level language tasks . Grefenstette et al . ( 2015 ) ; Dyer et al . ( 2015 ) ; Joulin & Mikolov ( 2015 ) augment RNNs with a stack structure that works as a natural parsing tool , and use them to process algorithmic and nature language processing ( NLP ) tasks ; nonetheless , the speed of stack-augmented RNNs is rather slow due to multiple push-pop operations at every time-step . Ke et al . ( 2018 ) propose the Sparse Attentive Backtracking ( SAB ) architecture , which recalls a few past hidden sates at every time-step and do “ mental ” backpropogations to the nearby hidden states with respect to the past hidden states . Gulcehre et al . ( 2017 ) propose the TARDIS network , which is an LSTM-resembled RNN that directly stores a fixed number of previous hidden states with learned key parameters for memory addressing . At every time-step , the network reads out one historical state rt from the memory , and uses it , along with the input xt and last hidden state ht−1 , to produce new hidden state ht , then overwrites ht to the location of rt . TARDIS optionally uses the gumbel-softmax estimator ( Maddison et al. , 2016 ; Jang et al. , 2016 ) to sample the location of rt ( we will explain this in section 2.1 ) , so the whole network is differentiable . However , the TARDIS still involves some hand-crafted memory addressing method and its addressing mechanism requires a considerable amount of memory and time consumption . Inspired by the TARDIS , we introduce the ARMIN architecture , a more light-weight and memorybased architecture with a very simple memory addressing mechanism and in-depth modification of the existing LSTM structure . Concretely , our contributions are as follows : • We distill previous memory addressing methods and propose a simple yet effective memory addressing mechanism for the external memory , namely Auto-addressing , by encoding the information for memory addressing directly via the inputs xt and the hidden states ht−1 . • We propose a novel recurrent cell that combines the gating advantages of LSTM and allows direct gradient backpropagation from output to memory . With only a single cell , it learns better representations than many hierarchical RNN structures . • We show in our ablation study that the ARMIN is robust to small iteration lengths when training long sequential data , which enables training with large batch sizes and boost its training speed . We further demonstrate in char-level language modelling task that the ARMIN obtain 40 % training speed gain than LSTM while keeping similar performance and memory consumption . • We demonstrate competitive results on various tasks while keeping efficient time and memory consumption during training and inference time . 2 BACKGROUND . 2.1 GUMBEL-SOFTMAX ESTIMATOR . Categorical distribution is a natural choice for representing discrete structure in the world ; however , it ’ s rarely used in neural networks due to its inability to backpropagate through samples ( Jang et al. , 2016 ) . To this end , Maddison et al . ( 2016 ) and Jang et al . ( 2016 ) propose a continuous relaxation of categorical distribution and the corresponding gumbel-softmax gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample . Specifically , given a probability distribution p = ( π1 , π2 , ... , πk ) over k categories , the gumbel-softmax estimator produces an one-hot sample vector y with its i-th element calculated as follows : yi = exp ( ( log ( πi ) + gi ) /τ ) ∑k j=1 exp ( ( log ( πj ) + gj ) /τ ) for i = 1 , 2 , ... , k , ( 1 ) where g1 , ... , gk are i.i.d samples drawn from Gumbel distribution ( Gumbel , 1954 ) : gi = − log ( − log ( ui ) ) , ui ∼ Uniform ( 0 , 1 ) , ( 2 ) and τ is the temperature parameter . In practice , we usually start at a high temperature and anneal to a small but non-zero temperature ( Jang et al. , 2016 ) . 3 AUTO-ADDRESSING AND RECURRENT MEMORY INTEGRATING NETWORK . In this section , we describe the structure of the ARMIN network as shown in Figure 1 . It consists of a recurrent cell and a small external memory that stores historical hidden states . While processing sequential data , the ARMIN performs reading from memory , cell processing and writing to memory operations in chronological order during each time-step . In the following subsections , we first explain the structure of the recurrent cell , and then discuss the read and write operations . 3.1 THE RECURRENT CELL OF ARMIN . Inspired by classical LSTM structure ( please refer to APPENDIX A for details in RNN and LSTM structures ) , we propose a novel memory-augmented recurrent cell structure , namely the ARMIN cell . At every time-step , it takes in an input xt , the last hidden state ht−1 and a recovered historical hidden state rt chosen by a read operation , and produces an output vector ot and the new hidden state ht . The computation process is as follows : { ght grt } = { σ σ } Wig [ xt , ht−1 , rt ] + big , ( 3 ) hgt−1 = g h t ◦ ht−1 , ( 4 ) rgt = g r t ◦ rt , ( 5 ) it ft gt oht ort = σ σ tanh σ σ Wgo [ xt , h g t−1 , r g t ] + bgo , ( 6 ) ht = ft ◦ ht−1 + it ◦ gt , ( 7 ) ot = [ o h t ◦ tanh ( ht ) , ort ◦ tanh ( rt ) ] . ( 8 ) where ht−1 , ht ∈ Rdh , rt ∈ Rdr , xt ∈ Rdi , and Wig ∈ R2dh× ( di+dh+dr ) , Wgo ∈ R ( 4dh+dr ) × ( di+dh+dr ) , big ∈ Rdh+dr , bgo ∈ R4dh+dr , ot ∈ Rdh+dr . We refer to dh as the hidden size of the recurrent cell of ARMIN . Usually we have dh = dr and allocate equal number of weight parameters for ht and rt . In equation 3 ∼ 5 , two gates are calculated to control the information flow for ht−1 and rt respectively , generating gated hidden state hgt−1 and historical state r g t ; using this method , we can filter out the irrelevant information for the current time-step . Then as shown in equation 6 , we compute the input gate it , forget gate ft , cell state gt and output gate oht for the new hidden state just like in classical LSTM structure . Additionally , an output gate ort for historical state rt is computed . Next in equation 7 , we compute new hidden state ht that is the sum of ht−1 and cell state gt , leveraged by forget gate ft and input gate it . Finally in equation 8 , we calculate the output of this timestep , which is the concatenation of the gated contents from ht and rt . Using this method , we can selectively backpropagate gradients from the output to the historical state rt . We involve rt ( or rgt ) in all gates ’ computations . Intuitively , the ht−1 acts as the old working memory , and rt is treated as the long-term memory . The cell processes them with the input xt to generate the new working memory ht and the output ot . More specifically , each rt is a summary of historical hidden states selected by auto-addressing mechanism . The ARMIN cell learns to recurrently integrate the summary of long-term information from rt into the working memory ht . We will demonstrate the efficiency of recurrent memory integration in language modelling tasks . Unlike the LSTM that passes a tuple of ( ht , ct ) to the next recurrence and outputs ht at the same time , the ARMIN only passes one hidden state ht and separately uses ot as output , which has 2 benefits : one is that it takes lower memory consumption to store one hidden state rather than two , and the other is we can decouple the information needed for the output ot , from the information ( i.e . ht ) that is needed for memorizing and being used at later time-steps . Furthermore , we list 3 important differences between the ARMIN cell and TARDIS cell , please refer to APPENDIX B.1 .

Summary: This paper introduces a memory-augmented RNN (MARNN) which aims at being lightweight and   differentiable. In a nutshell, authors propose to augment a LSTM-type architecture with several memory cells. At each time-step, MARNN retrieves one memory cell, updates his state, and updates the memory cell content. To learn the retrieval operation that requires discrete addressing,  authors rely on the Gumbel-Softmax. Authors evaluate their approach on PennTreeBank character level modelling where they demonstrate competitive performances. They also report state-of-art performance on the Thumos dataset. The paper is overall clear and pleasant to read.