Learning to Remember More  with Less Memorization
Year: 2019

Abstract:
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.

text: 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.

heading: 1 Introduction
text: 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.

heading: 2 Methods

heading: 2.1 Theoretical Analysis
text: 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.

heading: 2.1.1 Memory analysis of RNNs
text: 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.

heading: 2.1.2 Memory analysis of MANNs
text: 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.

heading: 2.2 Proposed Models
text: Uniform writing can apply to any MANNs that support writing operations. Since the writing intervals are discrete, i.e., li ∈ N+, UW is implemented as the following:
Mt =
{ fw (ot,Mt−1) if t = ⌊ T
D+1
⌋ k, k ∈ N+
Mt−1 otherwise (8)
By following Eq. (8), the write intervals are close to the optimal interval defined in Theorem 3 and approximately maximize the average contribution. This writing policy works well if timesteps are equally important and the task is to remember all of them to produce outputs (i.e., in copy task). However, in reality, timesteps are not created equal and a good model may need to ignore unimportant or noisy timesteps. That is why overwriting in MANN can be necessary. In the next section, we propose a method that tries to balance between following the optimal strategy and employing overwriting mechanism as in current MANNs.

heading: 2.2.1 Local optimal design
text: To relax the assumptions of Theorem 3, we propose two improvements of the Uniform Writing (UW) strategy. First, the intervals between writes are equal with length L (1 ≤ L ≤ ⌊ T
D+1
⌋ ). If L = 1, the strategy becomes regular writing and if L = ⌊ T
D+1
⌋ , it becomes
uniform writing. This ensures that after ⌊ T L ⌋ writes, all memory slots should be filled and the model has to learn to overwrite. Meanwhile, the average kept information is still locally maximized every L ∗D timesteps. Second, we introduce a cache of size L to store the hidden states of the controller during a write interval. Instead of using the hidden state at the writing timestep to update the memory, we perform an attention over the cache to choose the best representative hidden state. The model will learn to assign attention weights to the elements in the cache. This mechanism helps the model consider the importance of each timestep input in the local interval and thus relax the equal contribution assumption of Theorem 3. We name the writing strategy that uses the two mentioned-above improvements as Cached Uniform Writing (CUW). An illustration of the writing mechanism is depicted in Fig. 1.

heading: 2.2.2 Local memory-augmented attention unit
text: In this subsection, we provide details of the attention mechanism used in our CUW. To be specific, the best representative hidden state at is computed as follows:
αtj = softmax ( vT tanh (Wht−1 + Udj + V rt−1) ) (9) at = L∑ j=1 αtjdj (10)
where αtj is the attention score between the t-th writing step and the j-th element in the cache; W , U , V and v are parameters; h and r are the hidden state of the controller and the read-out (Eq. (1)), respectively; dj is the cache element and can be implemented as the controller’s hidden state (dj = ht−1−L+j).
The vector at will be used to replace the previous hidden state in updating the controller and memory. The whole process of performing CUW is summarized in Algo. 1.

heading: 3 Results

heading: 3.1 An Ablation Study: Memory-augmented Neural Networks with and without Uniform Writing
text: In this section, we study the impact of uniform writing on MANNs under various circumstances (different controller types, memory types and number of memory slots). We restrict the memorization problem to the double task in which the models must reconstruct a sequence of integers sampled uniformly from range [1, 10] twice. We cast this problem to a sequence to sequence problem with 10 possible outputs per decoding step. The training
stops after 10,000 iterations of batch size 64. We choose DNC1 and NTM2 as the two MANNs in the experiment. The recurrent controllers can be RNN or LSTM. With LSTM controller, the sequence length is set to 50. We choose sequence length of 30 to make it easier for the RNN controller to learn the task. The number of memory slots D is chosen from the set {2, 4, 9, 24} and {2, 4, 9, 14} for LSTM and RNN controllers, respectively. More memory slots will make UW equivalent to the regular writing scheme. For this experiment, we use Adam optimizer (Kingma & Ba, 2014) with initial learning rate and gradient clipping of {0.001, 0.0001} and {1, 5, 10}, respectively. The metric used to measure the performance is the average accuracy across decoding steps. For each configuration of hyper-parameters, we run the experiment 5 times and report the mean accuracy with error bars.
Figs. 2(a) and (c) depict the performance of UW and regular writing under different configurations. In any case, UW boosts the prediction accuracy of MANNs. The performance gain can be seen clearly when the compression ratio is between 10− 40%. This is expected since when the compression ratio is too small or too big, UW converges to regular writing. Interestingly, increasing the memory size does not always improve the performance, as in the case of NTM with RNN controllers. Perhaps, learning to attend to many memory slots is tricky for some task given limited amount of training data. This supports the need to apply UW to MANN with moderate memory size. We also conduct experiments to verify the benefit of using UW for bigger memory. The results can be found in Appendix H.
We also measure the speed-up of training time when applying UW on DNC and NTM, which is illustrated in Figs. 2(b) and (d). The result shows that with UW, the training time can drop up to 60% for DNC and 28% for NTM, respectively. As DNC is more complicated than NTM, using UW to reduce memory access demonstrates clearer speed-up in training (similar behavior can be found for testing time).

heading: 3.2 Synthetic Memorization
text: Here we address a broader range of baselines on two synthetic memorization tasks, which are the sequence copy and reverse. In these tasks, there is no discrimination amongst timesteps so the model’s goal is to learn to compress the input efficiently for later retrieval. We experiment with different sequence lengths of 50 and 100 timesteps. Other details are the same as the previous double task except that we fix the learning rate and gradient clipping to 0.001 and 10, respectively. The standard baselines include LSTM, NTM and DNC. All memory-augmented models have the same memory size of 4 slots, corresponding to compression ratio of 10% and 5%, respectively. We aim at this range of compression ratio to match harsh practical requirements. UW and CUW (cache size L = 5) are built upon the DNC, which from our previous observations, works best for given compression ratios. We choose different dimensions Nh for the hidden vector of the controllers to ensure the model
1Our reimplementation based on https://github.com/deepmind/dnc 2https://github.com/MarkPKCollier/NeuralTuringMachine
sizes are approximately equivalent. To further verify that our UW is actually the optimal writing strategy, we design a new baseline, which is DNC with random irregular writing strategy (RW). The write is sampled from a binomial distribution with p = (D + 1) /T (equivalent to compression ratio). After sampling, we conduct the training for that policy. The final performances of RW are taken average from 3 different random policies’ results.
The performance of the models is listed in Table 1. As clearly seen, UW is the best performer for the pure memorization tests. This is expected from the theory as all timesteps are importantly equivalent. Local attention mechanism in CUW does not help much in this scenario and thus CUW finishes the task as the runner-up. Reverse seems to be easier than copy as the models tend to “remember” more the last-seen timesteps whose contributions λT−t remains significant. In both cases, other baselines including random irregular and regular writing underperform our proposed models by a huge margin.

heading: 3.3 Synthetic Reasoning
text: Tasks in the real world rarely involve just memorization. Rather, they require the ability to selectively remember the input data and synthesize intermediate computations. To investigate whether our proposed writing schemes help the memory-augmented models handle these challenges, we conduct synthetic reasoning experiments which include add and max tasks. In these tasks, each number in the output sequence is the sum or the maximum of two numbers in the input sequence. The pairing is fixed as: yt = xt+xT−t 2 , t = 1, ⌊ T 2 ⌋ for
add task and yt = max (x2t, x2t+1) , t = 1, ⌊ T 2 ⌋ for max task, respectively. The length of the output sequence is thus half of the input sequence. A brief overview of input/output format for these tasks can be found in Appendix G. We deliberately use local (max) and distant (add) pairing rules to test the model under different reasoning strategies. The same experimental setting as in the previous section is applied except for the data sample range for the max task, which is [1, 50]3. LSTM and NTM are excluded from the baselines as they fail on these tasks.
Table 2 shows the testing results for the reasoning tasks. Since the memory size is small compared to the number of events, regular writing or random irregular writing cannot compete with the uniform-based writing policies. Amongst all baselines, CUW demonstrates superior performance in both tasks thanks to its local attention mechanism. It should be noted that the timesteps should not be treated equally in these reasoning tasks. The model should weight a timestep differently based on either its content (max task) or location (add task) and maintain its memory for a long time by following uniform criteria. CUW is designed to balance the two approaches and thus it achieves better performance. Further insights into memory operations of these models are given in Appendix I.

heading: 3.4 Synthetic Sinusoidal Regression
text: In real-world settings, sometimes a long sequence can be captured and fully reconstructed by memorizing some of its feature points. For examples, a periodic function such as sinusoid can be well-captured if we remember the peaks of the signal. By observing the peaks,
3With small range like [1, 10], there is no much difference in performance amongst models
we can deduce the frequency, amplitude, phase and thus fully reconstructing the function. To demonstrate that UW and CUW are useful for such scenarios, we design a sequential continuation task, in which the input is a sequence of sampling points across some sinusoid: y = 5+A sin(2πfx+ϕ). Here, A ∼ U (1, 5), f ∼ U (10, 30) and ϕ ∼ U (0, 100). After reading the input y = {yt}Tt=1, the model have to generate a sequence of the following points in the sinusoid. To ensure the sequence y varies and covers at least one period of the sinusoid, we set x = {xt}Tt=1 where xi = (t+ 1) /1000, 1 ∼ U (−1, 1). The sequence length for both input and output is fixed to T = 100. The experimental models are LSTM, DNC, UW and CUW (built upon DNC). For each model, optimal hyperparameters including learning rate and clipping size are tuned with 10,000 generated sinusoids. The memories have 4 slots and all baselines have similar parameter size. We also conduct the experiment with noisy inputs by adding a noise 2 ∼ U (−2, 2) to the input sequence y. This increases the difficulty of the task. The loss is the average of mean square error (MSE) over decoding timesteps.
We plot the mean learning curves with error bars over 5 runnings for sinusoidal regression task under clean and noisy condition in Figs. 3(a) and (b), respectively. Regular writing DNC learns fast at the beginning, yet soon saturates and approaches the performance of LSTM (MSE = 1.05 and 1.39 in clean and noisy condition, respectively). DNC performance does not improve much as we increase the memory size to 50, which implies the difficulty in learning with big memory. Although UW starts slower, it ends up with lower errors than DNC and perform slightly better than CUW in clean condition (MSE = 0.44 for UW and 0.61 for CUW). CUW demonstrates competitive performance against other baselines, approaching to better solution than UW for noisy task where the model should discriminate the timesteps (MSE = 0.98 for UW and 0.55 for CUW). More visualizations can be found in Appendix J.

heading: 3.5 Flatten Image Recognition
text: We want to compare our proposed models with DNC and other methods designed to help recurrent networks learn longer sequence. The chosen benchmark is a pixel-by-pixel image classification task on MNIST in which pixels of each image are fed into a recurrent model sequentially before a prediction is made. In this task, the sequence length is fixed to 768
with highly redundant timesteps (black pixels). The training, validation and testing sizes are 50,000, 10,000 and 10,000, respectively. We test our models on both versions of nonpermutation (MNIST) and permutation (pMNIST). More details on the task and data can be found in (Le et al., 2015). For DNC, we try with several memory slots from {15, 30, 60} and report the best results. For UW and CUW, memory size is fixed to 15 and cache size L is set to 10. The controllers are implemented as single layer GRU with 100-dimensional hidden vector. To optimize the models, we use RMSprop with initial learning rate of 0.0001.
Table 3 shows that DNC underperforms r-LSTM, which indicates that regular DNC with big memory finds it hard to beat LSTM-based methods. After applying UW, the results get better and with CUW, it shows significant improvement over r-LSTM and demonstrates competitive performance against dilated-RNNs models. Notably, dilated-RNNs use 9 layers in their experiments compared to our singer layer controller. Furthermore, our models exhibit more consistent performance than dilated-RNNs. For completeness, we include comparisons between CUW and non-recurrent methods in Appendix K

heading: 3.6 Document Classification
text: To verify our proposed models in real-world applications, we conduct experiments on document classification task. In the task, the input is a sequence of words and the output is the classification label. Following common practices in (Yogatama et al., 2017; Seo et al., 2018), each word in the document is embedded into a 300-dimensional vector using Glove embedding (Pennington et al., 2014). We use RMSprop for optimization, with initial learning rate of 0.0001. Early-stop training is applied if there is no improvement after 5 epochs in the validation set. Our UW and CUW are built upon DNC with single layer 512-dimensional LSTM controller and the memory size is chosen in accordance with the average length of the document, which ensures 10 − 20% compression ratio. The cache size for CUW is fixed to 10. The datasets used in this experiment are common big datasets where the number of documents is between 120,000 and 1,400,000 with maximum of 4,392 words per document (see Appendix L for further details). The baselines are recent state-of-the-arts in the domain, some of which are based on recurrent networks such as D-LSTM (Yogatama et al., 2017) and Skim-LSTM (Seo et al., 2018). We exclude DNC from the baselines as it is inefficient to train the model with big document datasets.
Our results are reported in Table 4. On five datasets out of six, our models beat or match the best published results. For IMDb dataset, our methods outperform the best recurrent model (Skim-LSTM). The performance gain is competitive against that of the state-of-thearts. In most cases, CUW is better than UW, which emphasizes the importance of relaxing the timestep equality assumption in practical situations. Details results across different runs for our methods are listed in Appendix M.
4Methods that use semi-supervised training to achieve higher accuracy are not listed.

heading: 4 Related Work
text: Traditional recurrent models such as RNN/LSTM (Elman, 1990; Hochreiter & Schmidhuber, 1997) exhibit some weakness that prevent them from learning really long sequences. The reason is mainly due to the vanishing gradient problem (Pascanu et al., 2013) or to be more specific, the exponential decay of input value over time. One way to overcome this problem is enforcing the exponential decay factor close to one by putting a unitary constraint on the recurrent weight (Arjovsky et al., 2016; Wisdom et al., 2016). Although this approach is theoretically motivated, it restricts the space of learnt parameters.
More relevant to our work, the idea of using less or adaptive computation for good has been proposed in (Graves, 2016; Yu et al., 2017; 2018; Seo et al., 2018). Most of these works are based on the assumption that some of timesteps in a sequence are unimportant and thus can be ignored to reduce the cost of computation and increase the performance of recurrent networks. Different form our approach, these methods lack theoretical supports and do not directly aim to solve the problem of memorizing long-term dependencies.
Dilated RNN (Chang et al., 2017) is another RNN-based proposal which improves long-term learning by stacking multiple dilated recurrent layers with hierarchical skip-connections. This theoretically guarantees the mean recurrent length and shares with our method the idea to construct a measurement on memorization capacity of the system and propose solutions to optimize it. The difference is that our system is memory-augmented neural networks while theirs is multi-layer RNNs, which leads to totally different optimization problems.
Recent researches recommend to replace traditional recurrent models by other neural architectures to overcome the vanishing gradient problem. The Transformer (Vaswani et al., 2017) attends to all timesteps at once, which ensures instant access to distant timestep yet requires quadratic computation and physical memory proportional to the sequence length. Memoryaugmented neural networks (MANNs), on the other hand, learn to establish a limited-size memory and attend to the memory only, which is scalable to any-length sequence. Compared to others, MANNs resemble both computer architecture design and human working memory (Logie, 2014). However, the current understanding of the underlying mechanisms and theoretical foundations for MANN are still limited.
Recent works on MANN rely almost on reasonable intuitions. Some introduce new addressing mechanisms such as location-based (Graves et al., 2014), least-used (Santoro et al., 2016) and order-based (Graves et al., 2016). Others focus on the scalability of MANN by using sparse memory access to avoid attending to a large number of memory slots (Rae et al., 2016). These problems are different from ours which involves MANN memorization capacity optimization.
Our local optimal solution to this problem is related to some known neural caching (Grave et al., 2017b;a; Yogatama et al., 2018) in terms of storing recent hidden states for later encoding uses. These methods either aim to create structural bias to ease the learning process (Yogatama et al., 2018) or support large scale retrieval (Grave et al., 2017a). These are different from our caching purpose, which encourages overwriting and relaxes the equal
contribution assumption of the optimal solution. Also, the details of implementation are different as ours uses local memory-augmented attention mechanisms.

heading: 5 Conclusions
text: We have introduced Uniform Writing (UW) and Cached Uniform Writing (CUW) as faster solutions for longer-term memorization in MANNs. With a comprehensive suite of synthetic and practical experiments, we provide strong evidences that our simple writing mechanisms are crucial to MANNs to reduce computation complexity and achieve competitive performance in sequence modeling tasks. In complement to the experimental results, we have proposed a meaningful measurement on MANN memory capacity and provided theoretical analysis showing the optimality of our methods. Further investigations to tighten the measurement bound will be the focus of our future work.

heading: A Derivation on the bound inequality in linear dynamic system
text: The linear dynamic system hidden state is described by the following recursive equation:
ht =Wxt + Uht−1 + b
By induction,
ht = t∑ i=1 U t−iWxi + C
where C is some constant with respect to xi. In this case, ∂ht∂xi = U t−iW . By applying norm sub-multiplicativity5,
ci−1,t = ∥∥U t−i+1W∥∥
≤ ‖U‖ ∥∥U t−iW∥∥
= ‖U‖ ci,t
That is, λc = ‖U‖.

heading: B Derivation on the bound inequality in standard RNN
text: The standard RNN hidden state is described by the following recursive equation:
ht = tanh (Wxt + Uht−1 + b)
From ∂ht∂xi = ∂ht ∂ht−1 ∂ht−1 ∂xi , by induction,
∂ht ∂xi =  t∏ j=i+1 ∂hj ∂hj−1  ∂hi ∂xi =  t∏ j=i diag ( tanh′ (aj) )U t−iW where aj = Wxj + Uhj−1 + b and diag (·) converts a vector into a diagonal matrix. As 0 ≤ tanh′ (x) = 1 − tanh (x)2 ≤ 1,
∥∥diag (tanh′ (X))∥∥ is bounded by some value B. By applying norm sub-multiplicativity,
ci−1,t = ∥∥∥∥∥∥  t∏ j=i diag ( tanh′ (aj) ) diag (tanh′ (ai−1))U t−i+1W ∥∥∥∥∥∥
≤ ‖U‖ ∥∥∥∥∥∥ t∏ j=i diag ( tanh′ (aj) ) U t−iW ∥∥∥∥∥∥∥∥diag (tanh′ (ai−1))∥∥ = B ‖U‖ ci,t
That is, λc = B ‖U‖.
5If not explicitly stated otherwise, norm refers to any consistent matrix norm which satisfies sub-multiplicativity.

heading: C Derivation on the bound inequality in LSTM
text: For the case of LSTM, the recursive equation reads:
ct = σ (Ufxt +Wfht−1 + bf ) ct−1 + σ (Uixt +Wiht−1 + bi) tanh (Uzxt +Wzht−1 + bz)
ht = σ (Uoxt +Woht−1 + bo) tanh (ct)
Taking derivatives,
∂hj ∂hj−1 = σ′ (oj) tanh (cj)Wo + σ (oj) tanh ′ (cj)σ ′ (fj) cj−1Wf
+ σ (oj) tanh ′ (cj)σ ′ (ij) tanh (zj)Wi + σ (oj) tanh ′ (cj)σ (ij) tanh ′ (zj)Wz
∂hj−1 ∂xj−1 = σ′ (oj−1) tanh (cj−1)Uo + σ (oj−1) tanh ′ (cj−1)σ ′ (fj−1) cj−2Uf
+ σ (oj−1) tanh ′ (cj−1)σ ′ (ij−1) tanh (zj−1)Ui + σ (oj−1) tanh ′ (cj−1)σ (ij−1) tanh ′ (zj−1)Uz
∂hj ∂xj = σ′ (oj) tanh (cj)Uo + σ (oj) tanh ′ (cj)σ ′ (fj) cj−1Uf
+ σ (oj) tanh ′ (cj)σ ′ (ij) tanh (zj)Ui + σ (oj) tanh ′ (cj)σ (ij) tanh ′ (zj)Uz
where oj denotes the value in the output gate at j-th timestep (similar notations are used for input gate (ij), forget gate (fj) and cell value (zj)) and “non-matrix” terms actually represent diagonal matrices corresponding to these terms. Under the assumption that h0=0, we then make use of the results in (Miller & Hardt, 2018) stating that ‖ct‖∞ is bounded for all t. By applying l∞-norm sub-multiplicativity and triangle inequality, we can show that
∂hj ∂hj−1 ∂hj−1 ∂xj−1 =M ∂hj ∂xj +N
with
‖M‖∞ ≤ 1/4 ‖Wo‖∞ + 1/4 ‖Wf‖∞ ‖cj‖∞ + 1/4 ‖Wi‖∞ + ‖Wz‖∞ = Bm ‖N‖∞ ≤ 1/16 ‖WoUi‖∞ + 1/16 ‖WoUf‖∞ ( ‖cj‖∞ + ‖cj−1‖∞ ) + 1/4 ‖WoUz‖∞
+ 1/16 ‖WiUo‖∞ + 1/16 ‖WiUf‖∞ ( ‖cj‖∞ + ‖cj−1‖∞ ) + 1/4 ‖WiUz‖∞
+ 1/16 ( ‖WfUo‖∞ + 1/16 ‖WfUi‖∞ + 1/4 ‖WfUz‖∞ ) ( ‖cj‖∞ + ‖cj−1‖∞ ) + 1/4 ‖WzUo‖∞ + 1/4 ‖WzUf‖∞ ( ‖cj‖∞ + ‖cj−1‖∞ ) + 1/4 ‖WzUi‖∞
= Bn
By applying l∞-norm sub-multiplicativity and triangle inequality,
ci−1,t = ∥∥∥∥∥∥ t∏
j=i+1
∂hj ∂hj−1 ∂hi ∂hi−1 ∂hi−1 ∂xi−1 ∥∥∥∥∥∥ ∞
= ∥∥∥∥∥∥ t∏
j=i+1
∂hj ∂hj−1 ( ∂hi ∂xi m+ n )∥∥∥∥∥∥ ∞
≤ Bmci,t +Bn t∏
j=i+1
∥∥∥∥ ∂hj∂hj−1 ∥∥∥∥ ∞
As LSTM is λ-contractive with λ < 1 in the l∞-norm (readers are recommended to refer to (Miller & Hardt, 2018) for proof), which implies ∥∥∥ ∂hj∂hj−1 ∥∥∥∞ < 1, Bn t∏j=i+1 ∥∥∥ ∂hj∂hj−1 ∥∥∥∞ → 0 as t− i→∞. For t− i <∞, under the assumption that ∂hj∂xj 6= 0, we can always find some value B <∞ such that ci−1,t ≤ Bci,t. For t−i→∞, λc → Bm. That is, λc = max (Bm, B).

heading: D Proof of theorem 1
text: Proof. Given that λcci,t ≥ ci−1,t with some λc ∈ R+, we can use ct,tλt−ic as the upper bound on ci,t with i = 1, t, respectively. Therefore,
f(0) ≤ T∑ t=1 ct,T ≤ cT,T T∑ t=1 λT−tc = f (λc)
where f(λ) = cT,T T∑ t=1 λT−t is continuous on R+. According to intermediate value theorem,
there exists λ ∈ (0, λc] such that cT,T T∑ t=1 λT−t = T∑ t=1 ct,T .

heading: E Proof of theorem 2
text: Proof. According to Theorem 1, there exists some λi ∈ R+such that the summation of contribution stored between Ki and Ki+1 can be quantified as cKi+1,Ki+1 Ki+1∑ t=Ki λ Ki+1−t i
(after ignoring contributions before Ki-th timestep for simplicity). Let denote P (λ) = Ki+1∑ t=Ki
λKi+1−t, we have P ′ (λ) > 0,∀λ ∈ R+. Therefore, P (λi) ≥ P ( min i (λi) ) . Let C = min i (ci,i) and λ = min i (λi), the average contribution stored in a MANN has a lower bound quantified
as Iλ, where 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 .

heading: F Proof of theorem 3
text: Proof. The second-order derivative of fλ (x) reads:
f ′′λ (x) = − (lnλ)
2
1− λ λx (11)
We have f ′′λ, (x) ≤ 0 with ∀x ∈ R+ and 1 > λ > 0, so fλ (x) is a concave function. Thus, we can apply Jensen inequality as follows:
1
D + 1 D+1∑ i=1 fλ(li) ≤ fλ ( D+1∑ i=1 1 D + 1 li ) = fλ ( T D + 1 ) (12)
Equality holds if and only if l1 = l2 = ... = lD+1 = TD+1 . We refer to this as Uniform Writing strategy. By plugging the optimal values of li, we can derive the maximized average contribution as follows:
Iλmax ≡ gλ (T,D) = C (D + 1)
T
( 1− λ T D+1
1− λ
) (13)
When λ = 1, Iλ = CT D+1∑ i=1 li = C. This is true for all writing strategies. Thus, Uniform Writing is optimal for 0 < λ ≤ 1.
We can show that this solution is also optimal for the case λ > 1. As f ′′λ (x) > 0 with ∀x ∈ R+;λ > 1, fλ (x) is a convex function and Eq. (12) flips the inequality sign. Thus, Iλ reaches its minimum with Uniform Writing. For λ > 1, minimizing Iλ is desirable to prevent the system from diverging.
We can derive some properties of function g. Let x = D+1L , gλ (L,D) = gλ (x) = Cx( λ 1 x−1 λ−1 ). We have g′λ (x) = Cλ ( 1− λ 1x ) (x− lnλ) > 0 with 0 < λ ≤ 1,∀x ≥ 0, so gλ (T,D) is an
increasing function if we fix T and let D vary. That explains why having more memory slots helps improve memorization capacity. If D = 0, gλ (T, 0) becomes E.q (5). In this case, MANNs memorization capacity converges to that of recurrent networks.

heading: G Summary of synthetic discrete task format

heading: H UW performance on bigger memory

heading: I Memory operating behaviors on synthetic tasks
text: In this section, we pick three models (DNC, DNC+UW and DNC+CUW) to analyze their memory operating behaviors. Fig. 4 visualizes the values of the write weights and read weights for the copy task during encoding input and decoding output sequence, respectively. In the copy task, as the sequence length is 50 while the memory size is 4, one memory slot should contain the accumulation of multiple timesteps. This principle is reflected in the decoding process in three models, in which one memory slot is read repeatedly across several timesteps. Notably, the number of timesteps consecutively spent for one slot is close to 10- the optimal interval, even for DNC ( Fig. 4(a)), which implies that the ultimate rule would be the uniform rule. As UW and CUW are equipped with uniform writing, their writing patterns follow the rule perfectly. Interestingly, UW chooses the first written location for the final write (corresponding to the <eos> token) while CUW picks the last written location. As indicated in Figs. 4(b) and (c), both of them can learn the corresponding reading pattern for decoding process, which leads to good performances. On the other hand, regular DNC fails to learn a perfect writing strategy. Except for the timesteps at the end of the sequence, the timesteps are distributed to several memory slots while the reading phase attends to one memory slot repeatedly. This explains why regular DNC cannot compete with the other two proposed methods in this task.
En co
di ng
W rit
e G
at e
For the max task, Fig. 5 displays similar visualization with an addition of write gate during encoding phase. The write gate indicates how much the model should write the input at some timestep to the memory. A zero write gate means there is no writing. For this task, a good model should discriminate between timesteps and prefer writing the greater ones. As clearly seen in Fig. 5(a), DNC suffers the same problem as in copy task, unable to synchronize encoding writing with decoding reading. Also, DNC’s write gate pattern does not show reasonable discrimination. For UW (Fig. 5(b)), it tends to write every timestep and relies on uniform writing principle to achieve write/read accordance and thus better results than DNC. Amongst all, CUW is able to ignore irrelevant timesteps and follows uniform writing at the same time (see Fig. 5(c)).
J Visualizations of model performance on sinusoidal regression tasks
We pick randomly 3 input sequences and plot the output sequences produced by DNC, UW and CUW in Figs. 6 (clean) and 7 (noisy). In each plot, the first and last 100 timesteps correspond to the given input and generated output, respectively. The ground truth sequence is plotted in red while the predicted in blue. We also visualize the values of MANN write gates through time in the bottom of each plots. In irregular writing encoding phase, the write gate is computed even when there is no write as it reflects how much weight the controller puts on the timesteps. In decoding, we let MANNs write to memory at every timestep to allow instant update of memory during inference.
Under clean condition, all models seem to attend more to late timesteps during encoding, which makes sense as focusing on late periods of sine wave is enough for later reconstruction. However, this pattern is not clear in DNC and UW as in CUW. During decoding, the write gates tend to oscillate in the shape of sine wave, which is also a good strategy as this directly reflects the amplitude of generation target. In this case, both UW and CUW demonstrate this behavior clearer than DNC.
Under noisy condition, DNC and CUW try to follow sine-shape writing strategy. However, only CUW can learn the pattern and assign write values in accordance with the signal period,
which helps CUW decoding achieve highest accuracy. On the other hand, UW choose to assign write value equally and relies only on its maximization of timestep contribution. Although it achieves better results than DNC, it underperforms CUW.
K Comparison with non-recurrent methods in flatten image classification task

heading: L Details on document classification datasets

heading: M Document classification detailed records

Review summaries:
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".

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.