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 looks at ways to improve memory-writing in memory augmented neural networks. Authors proposed two methods to compare against "regular writing" method as well as compare against each other, namely "uniform writing" and "cached uniform writing". Latter one attempts to utilize a small size memory efficiently by introducing memory overwriting in other words "forgetting".