Learning to Remember More with Less Memorization

Memory-augmented neural networks consisting of a neural controller and an external memory have shown potentials in long-term sequential learning . Current RAM-like memory models maintain memory accessing every timesteps , thus they do not effectively leverage the short-term memory held in the controller . We hypothesize that this scheme of writing is suboptimal in memory utilization and introduces redundant computation . To validate our hypothesis , we derive a theoretical bound on the amount of information stored in a RAM-like system and formulate an optimization problem that maximizes the bound . The proposed solution dubbed Uniform Writing is proved to be optimal under the assumption of equal timestep contributions . To relax this assumption , we introduce modifications to the original solution , resulting in a solution termed Cached Uniform Writing . This method aims to balance between maximizing memorization and forgetting via overwriting mechanisms . Through an extensive set of experiments , we empirically demonstrate the advantages of our solutions over other recurrent architectures , claiming the state-of-the-arts in various sequential modeling tasks . 1 Introduction . A core task in sequence learning is to capture long-term dependencies amongst timesteps which demands memorization of distant inputs . In recurrent neural networks ( RNNs ) , the memorization is implicitly executed via integrating the input history into the state of the networks . However , learning vanilla RNNs over long distance proves to be difficult due to the vanishing gradient problem ( Bengio et al. , 1994 ; Pascanu et al. , 2013 ) . One alleviation is to introduce skip-connections along the execution path , in the forms of dilated layers ( Van Den Oord et al. , 2016 ; Chang et al. , 2017 ) , attention mechanisms ( Bahdanau et al. , 2015 ; Vaswani et al. , 2017 ) and external memory ( Graves et al. , 2014 ; 2016 ) . Amongst all , using external memory most resembles human cognitive architecture where we perceive the world sequentially and make decision by consulting our memory . Recent attempts have simulated this process by using RAM-like memory architectures that store information into memory slots . Reading and writing are governed by neural controllers using attention mechanisms . These memory-augmented neural networks ( MANN ) have demonstrated superior performance over recurrent networks in various synthetic experiments ( Graves et al. , 2016 ) and realistic applications ( Le et al. , 2018a ; b ; Franke et al. , 2018 ) . Despite the promising empirical results , there is no theoretical analysis or clear understanding on optimal operations that a memory should have to maximize its performance . To the best of our knowledge , no solution has been proposed to help MANNs handle ultra-long sequences given limited memory . This scenario is practical because ( i ) sequences in the real-world can be very long while the computer resources are limited and ( ii ) it reflects the ability to compress in human brain to perform life-long learning . Previous attempts such as ( Rae et al. , 2016 ) try to learn ultra-long sequences by expanding the memory , which is not always feasible and do not aim to optimize the memory by some theoretical criterion . This paper presents a new approach towards finding optimal operations for MANNs that serve the purpose of learning longer sequences with finite memory . More specifically , upon analyzing RNN and MANN operations we first introduce a measurement on the amount of information that a MANN holds after encoding a sequence . This metric reflects the quality of memorization under the assumption that contributions from timesteps are equally important . We then derive a generic solution to optimize the measurement . We term this optimal solution as Uniform Writing ( UW ) , and it is applicable for any MANN due to its generality . Crucially , UW helps reduce significantly the computation time of MANN . Third , to relax the assumption and enable the method to work in realistic settings , we further propose Cached Uniform Writing ( CUW ) as an improvement over the Uniform Writing scheme . By combining uniform writing with local attention , CUW can learn to discriminate timesteps while maximizing local memorization . Finally we demonstrate that our proposed models outperform several MANNs and other state-of-the-art methods in various synthetic and practical sequence modeling tasks . 2 Methods . 2.1 Theoretical Analysis . Memory-augmented neural networks can be viewed as an extension of RNNs with external memory M . The memory supports read and write operations based on the output ot of the controller , which in turn is a function of current timestep input xt , previous hidden state ht−1 and read value rt−1 from the memory . Let assume we are given these operators from recent MANNs such as NTM ( Graves et al. , 2014 ) or DNC ( Graves et al. , 2016 ) , represented as : rt = fr ( ot , Mt−1 ) ( 1 ) Mt = fw ( ot , Mt−1 ) ( 2 ) The controller output and hidden state are updated as follows : ot = fo ( ht−1 , rt−1 , xt ) ( 3 ) ht = fh ( ht−1 , rt−1 , xt ) ( 4 ) Here , fo and fh are often implemented as RNNs while fr and fw are designed specifically for different memory types . Current MANNs only support regular writing by applying Eq . ( 2 ) every timestep . In effect , regular writing ignores the accumulated short-term memory stored in the controller hidden states which may well-capture the recent subsequence . We argue that the controller does not need to write to memory continuously as its hidden state also supports memorizing . Another problem of regular writing is time complexity . As the memory access is very expensive , reading/writing at every timestep makes MANNs much slower than RNNs . This motivates a irregular writing strategy to utilize the memorization capacity of the controller and consequently , speed up the model . In the next sections , we first define a metric to measure the memorization performance of RNNs , as well as MANNs . Then , we solve the problem of finding the best irregular writing that optimizes the metric . 2.1.1 Memory analysis of RNNs . We first define the ability to “ remember ” of recurrent neural networks , which is closely related to the vanishing/exploding gradient problem ( Pascanu et al. , 2013 ) . In RNNs , the state transition ht = φ ( ht−1 , xt ) contains contributions from not only xt , but also previous timesteps xi < t embedded in ht−1 . Thus , ht can be considered as a function of timestep inputs , i.e , ht = f ( x1 , x2 , ... , xt ) . One way to measure how much an input xi contributes to the value of ht is to calculate the norm of the gradient ∥∥∥∂ht∂xi ∥∥∥ . If the norm equals zero , ht is constant w.r.t xi , that is , ht does not “ remember ” xi . As a bigger ∥∥∥∂ht∂xi ∥∥∥ implies more influence of xi on ht , we propose using ∥∥∥∂ht∂xi ∥∥∥ to measure the contribution of the i-th input to the t-th hidden state . Let ci , t denotes this term , we can show that in the case of common RNNs , λcci , t ≥ ci−1 , t with some λc ∈ R+ ( see Appendix A - C for proof ) . This means further to the past , the contribution decays ( when λc < 1 ) or grows ( when λc > 1 ) with the rate of at least λc .We can measure the average amount of contributions across T timesteps as follows ( see Appendix D for proof ) : Theorem 1 . There exists λ ∈ R+such that the average contribution of a sequence of length T with respect to a RNN can be quantified as the following : Iλ = T∑ t=1 ct , T T = cT , T T∑ t=1 λT−t T ( 5 ) If λ < 1 , λT−t → 0 as T − t → ∞ . This is closely related to vanishing gradient problem . LSTM is known to “ remember ” long sequences better than RNN by using extra memory gating mechanisms , which help λ to get closer to 1 . If λ > 1 , the system may be unstable and suffer from the exploding gradient problem . 2.1.2 Memory analysis of MANNs . In slot-based MANNs , memory M is a set of D memory slots . A write at step t can be represented by the controller ’ s hidden state ht , which accumulates inputs over several timesteps ( i.e. , x1 , ... , xt ) . If another write happens at step t+k , the state ht+k ’ s information containing timesteps xt+1 , ... , xt+k is stored in the memory ( ht+k may involves timesteps further to the past , yet they are already stored in the previous write and can be ignored ) . During writing , overwriting may happen , replacing an old write with a new one . Thus after all , D memory slots associate with D chosen writes of the controller . From these observations , we can generalize Theorem 1 to the case of MANNs having D memory slots ( see Appendix E for proof ) . Theorem 2 . With any D chosen writes at timesteps 1 ≤K1 < K2 < ... < KD < T , there exist λ , C ∈ R+such that the lower bound on the average contribution of a sequence of length T with respect to a MANN having D memory slots can be quantified as the following : Algorithm 1 Cached Uniform Writing Require : a sequence x = { xt } Tt=1 , a cache C sized L , a memory sized D. 1 : for t = 1 , T do 2 : C.append ( ht−1 ) 3 : if t mod L == 0 then 4 : Use Eq . ( 10 ) to calculate at 5 : Execute Eq . ( 3 ) : ot = fo ( at , rt−1 , xt ) 6 : Execute Eq . ( 4 ) : ht = fh ( at , rt−1 , xt ) 7 : Update the memory using Eq . ( 2 ) 8 : Read rt from the memory using Eq . ( 1 ) 9 : C.clear ( ) 10 : else 11 : Update the controller using Eq . ( 4 ) : ht = fh ( ht−1 , rt−1 , xt ) 12 : Assign rt = rt−1 13 : end if 14 : end for Iλ = C K1∑ t=1 λK1−t+ K2∑ t=K1+1 λK2−t + ... + KD∑ t=KD−1+1 λKD−t+ T∑ t=KD+1 λT−t T = C T D+1∑ i=1 li−1∑ j=0 λj = C T D+1∑ i=1 fλ ( li ) ( 6 ) where li = K1 ; i = 1 Ki −Ki−1 ; D ≥ i > 1 T −KD ; i = D + 1 , fλ ( x ) = { 1−λx 1−λ λ 6= 1 x λ = 1 , ∀x ∈ R+ . If λ ≤ 1 , we want to maximize Iλ to keep the information from vanishing . On the contrary , if λ > 1 , we may want to minimize Iλ to prevent the information explosion . As both scenarios share the same solution ( see Appendix F ) , thereafter we assume that λ ≤ 1 holds for other analyses . By taking average over T , we are making an assumption that all timesteps are equally important . This helps simplify the measurement as Iλ is independent of the specific position of writing . Rather , it is a function of the interval lengths between the writes . This turns out to be an optimization problem whose solution is stated in the following theorem . Theorem 3 . Given D memory slots , a sequence with length T , a decay rate 0 < λ ≤ 1 , then the optimal intervals { li ∈ R+ } D+1 i=1 satisfying T = D+1∑ i=1 li such that the lower bound on the average contribution Iλ = CT D+1∑ i=1 fλ ( li ) is maximized are the following : l1 = l2 = ... = lD+1 = T D + 1 ( 7 ) We name the optimal solution as Uniform Writing ( UW ) and refer to the term TD+1 and D+1 T as the optimal interval and the compression ratio , respectively . The proof is given in Appendix F .

Summary: This paper deals with Memory Augmnted Neural Networks (MANN) and introduces an algorithm which allows full writes to the dense memory to be only exectued every L timesteps. The controller produces a hidden output at most timestps, whih is appended to a cache. Every L steps, soft attention is used to combine this cache of N hidden states to a single one, and then this is used as the input hidden state for the controller, with the outputs performing a write in the full memory M, along with clearing the cache.