Learning to Remember More with Less Memorization
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 eﬀectively 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 modiﬁcations 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.
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 diﬃcult 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 reﬂects 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 ﬁnding optimal operations for MANNs that serve the purpose of learning longer sequences with ﬁnite memory.
More speciﬁcally, upon analyzing RNN and MANN operations we ﬁrst introduce a measurement on the amount of information that a MANN holds after encoding a sequence.
This metric reﬂects 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 signiﬁcantly 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.
Methods 2. 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 o t of the controller, which in turn is a function of current timestep input x t , previous hidden state h t − 1 and read value r t − 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:
[Equation]
The controller output and hidden state are updated as follows:
[Equation]
Here, f o and f h are often implemented as RNNs while f r and f w are designed speciﬁcally for diﬀerent memory types.
Current MANNs only support regular writing by applying Eq.
(2) every timestep.
In eﬀect, 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 ﬁrst deﬁne a metric to measure the memorization performance of RNNs, as well as MANNs.
Then, we solve the problem of ﬁnding the best irregular writing that optimizes the metric.
2.1. Memory analysis of RNNs
We ﬁrst deﬁne 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 h t = φ ( h t − 1 , x t ) contains contributions from not only x t , but also previous timesteps x i<t embedded in h t − 1 .
Thus, h t can be considered as a function of timestep inputs, i.e, h t = f
( x 1 , x 2 , ..., x t ) .
One way to measure how much an input x i contributes to the value of h t is to calculate the norm of the gradient (cid:13)(cid:13)(cid:13) ∂h t ∂x i (cid:13)(cid:13)(cid:13) .
If the norm equals zero, h t is constant w.r.t x i , that is, h t does not “remember” x i .
As a bigger (cid:13)(cid:13)(cid:13) ∂h t ∂x i (cid:13)(cid:13)(cid:13) implies more inﬂuence of x i on h t , we propose using (cid:13)(cid:13)(cid:13) ∂h t ∂x i (cid:13)(cid:13)(cid:13) to measure the contribution of the i -th input [Figure] to the t -th hidden state.
Let c i,t denotes this term, we can show that in the case of common
RNNs, λ c c i,t ≥ c i − 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 quantiﬁed as the following:
[Equation]
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 suﬀer from the exploding gradient problem.
2.1. 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 h t , which accumulates inputs over several timesteps (i.e., x 1 , ..., x t ).
If another write happens at step t + k , the state h t + k ’s information containing timesteps x t +1 , ..., x t + k is stored in the memory ( h t + 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 ≤ K 1 < K 2 < ... <
K D < 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 quantiﬁed as the following:
Require: a sequence x = { x t } Tt =1 , a cache C sized L , a memory sized D .
1: for t = 1 , T do 2:
C .append
( h t − 1 ) 3: if t mod L == 0 then 4:
Use Eq.
( 10 ) to calculate a t 5:
Execute Eq.
( 3 ): o t = f
o ( a t , r t − 1 , x t ) 6:
Execute Eq.
( 4 ): h t = f h ( a t , r t − 1 , x t ) 7:
Update the memory using Eq.
Read r t from the memory using Eq.
C .clear() 10: else 11:
Update the controller using Eq.
( 4 ): h t = f h ( h t − 1 , r t − 1 , x t ) 12:
Assign r t = r t −13: end if 14: end for I λ = C K 1 (cid:80) t =1 λ K 1 − t +
K 2 (cid:80) t = K 1 + λ K 2 − t + ... +
K D (cid:80) t = K D − 1 + λ K D − t + T (cid:80) t = K D + λ T − t T = C T
D +1 (cid:88) i = l i − 1 (cid:88) j = λ j = C T D +1 (cid:88) i = f λ ( l i ) (6) where l i =  
K; i = K i − K
i −; D ≥ i >
T − K D ; i = D +, f λ ( x ) = (cid:40) 1 − λ x 1 − λ λ (cid:54) = 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 speciﬁc 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 { l i ∈ R + } D +1 i =1 satisfying T = D +1 (cid:80) i
=
1 l i such that the lower bound on the average contribution I λ =
CT D +1 (cid:80) i =
1 f λ ( l i ) is maximized are the following:
[Equation]
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.
2. Proposed Models
Uniform writing can apply to any MANNs that support writing operations.
Since the writing intervals are discrete, i.e., l i ∈ N + , UW is implemented as the following:
[Equation]
By following Eq.
(8), the write intervals are close to the optimal interval deﬁned 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.
2.2. Local optimal design
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 ≤ (cid:106) T D +1 (cid:107) ).
If L = 1 , the strategy becomes regular writing and if L = (cid:106) T D +1 (cid:107) , it becomes uniform writing.
This ensures that after (cid:4) TL (cid:5) writes, all memory slots should be ﬁlled 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.
2.2. Local memory-augmented attention unit
In this subsection, we provide details of the attention mechanism used in our CUW.
To be speciﬁc, the best representative hidden state a t is computed as follows:
[Equation]
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; d j is the cache element and can be implemented as the controller’s hidden state ( d j = h t − 1 − L + j ).
The vector a t 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.
Results 3. An Ablation Study: Memory-augmented Neural Networks with and without Uniform Writing
In this section, we study the impact of uniform writing on MANNs under various circumstances (diﬀerent 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 DNC 1 and NTM 2 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 conﬁguration 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 diﬀerent conﬁgurations.
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 beneﬁt 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).
3. Synthetic Memorization
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 eﬃciently for later retrieval.
We experiment with diﬀerent sequence lengths of 50 and 100 timesteps.
Other details are the same as the previous double task except that we ﬁx 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 diﬀerent dimensions N h for the hidden vector of the controllers to ensure the model [Table] 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 ﬁnal performances of RW are taken average from 3 diﬀerent 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 ﬁnishes 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 signiﬁcant.
In both cases, other baselines including random irregular and regular writing underperform our proposed models by a huge margin.
3. Synthetic Reasoning
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 ﬁxed as: y t
= x t + x T − t, t = 1 , (cid:4) T 2 (cid:5) for add task and y t = max ( x 2 t , x 2 t +1 ) , t = 1 , (cid:4) T 2 (cid:5) 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 diﬀerent 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 diﬀerently 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.
3. Synthetic Sinusoidal Regression
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, [Table] 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 = { y t } 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 = { x t } Tt =1 where x i =
( t + (cid:15) 1 ) / 1000 , (cid:15) 1 ∼ U ( − 1 , 1) .
The sequence length for both input and output is ﬁxed 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 (cid:15) 2 ∼ U ( − 2 , 2) to the input sequence y .
This increases the diﬃculty 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 .
in clean and noisy condition, respectively).
DNC performance does not improve much as we increase the memory size to 50 , which implies the diﬃculty 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.
3. Flatten Image Recognition
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 classiﬁcation 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 ﬁxed to
[Table]
Table 3 shows that DNC underperforms r-LSTM, which indicates that regular DNC with big memory ﬁnds it hard to beat LSTM-based methods.
After applying UW, the results get better and with CUW, it shows signiﬁcant 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
3. Document Classification
To verify our proposed models in real-world applications, we conduct experiments on document classiﬁcation task.
In the task, the input is a sequence of words and the output is the classiﬁcation 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 ﬁxed 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 ineﬃcient to train the model with big document datasets.
Our results are reported in Table 4.
On ﬁve 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 diﬀerent runs for our methods are listed in Appendix M. [Table]
Related Work
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 speciﬁc, 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.
Diﬀerent 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 diﬀerence is that our system is memory-augmented neural networks while theirs is multi-layer RNNs, which leads to totally diﬀerent 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 diﬀerent 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 diﬀerent from our caching purpose, which encourages overwriting and relaxes the equal contribution assumption of the optimal solution.
Also, the details of implementation are diﬀerent as ours uses local memory-augmented attention mechanisms.
Conclusions
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.
Appendix A Derivation on the bound inequality in linear dynamic system
The linear dynamic system hidden state is described by the following recursive equation:
[Equation]
By induction,
[Equation]
where C is some constant with respect to x i .
In this case, ∂h t ∂x i =
U t − i W .
By applying norm sub-multiplicativity 5 ,
[Equation]
That is, λ c = (cid:107) U (cid:107) .
B Derivation on the bound inequality in standard RNN
The standard RNN hidden state is described by the following recursive equation:
[Equation]
From ∂h t ∂x i =
∂h
t ∂h t − 1 ∂h t − 1 ∂x i , by induction, ∂h t ∂x
[Equation]
  t (cid:89) j
= i +
[Equation]
where a j = Wx j + Uh j − 1 + b and diag ( · ) converts a vector into a diagonal matrix.
As 0 ≤ tanh (cid:48) ( x ) = 1 − tanh ( x ) 2 ≤ 1 , (cid:13)(cid:13) diag (cid:0) tanh (cid:48) ( X ) (cid:1)(cid:13)(cid:13) is bounded by some value B .
By applying norm sub
-multiplicativity
[Equation]
That is, λ c = B (cid:107) U (cid:107)
[Equation]
C Derivation on the bound inequality in LSTM
For the case of LSTM, the recursive equation reads:
[Equation]
Taking derivatives,
[Equation]
where o j denotes the value in the output gate at j -th timestep (similar notations are used for input gate ( i j ), forget gate ( f j ) and cell value ( z j )) and “non-matrix” terms actually represent diagonal matrices corresponding to these terms.
Under the assumption that h
=
0, we then make use of the results in (Miller & Hardt, 2018) stating that (cid:107) c t (cid:107) ∞ is bounded for all t .
By applying l ∞ -norm sub-multiplicativity and triangle inequality
[Equation]
we can show that
[Equation]
with
[Equation]
= B n
By applying l ∞ -norm sub-multiplicativity and triangle inequality, c i − 1 ,t = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) t (cid:89)
j = i + ∂h
j ∂h j −∂h
i ∂h
i −∂h
i −∂x i −(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
∞
=
(cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) t (cid:89)
j = i + ∂h
j ∂h j −(cid:18)
∂h i ∂x i m + n (cid:19) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∞ ≤ B m c
i,t + B n t (cid:89) j
= i + (cid:13)(cid:13)(cid:13)(cid:13)
[Table]
As LSTM is λ -contractive with λ < 1 in the l ∞ -norm (readers are recommended to refer to (Miller & Hardt, 2018) for proof), which implies (cid:13)(cid:13)(cid:13)
∂h
j ∂h j − 1 (cid:13)(cid:13)(cid:13) ∞ < 1 , B n t (cid:81)
j = i + (cid:13)(cid:13)(cid:13)
[Table]
j [Table] − 1 (cid:13)(cid:13)(cid:13) ∞
→as t − i → ∞ .
For t − i < ∞ , under the assumption that ∂h j ∂x j (cid:54) = 0 , we can always ﬁnd some value B < ∞ such that c
i − 1 ,t ≤
Bc i,t .
For t − i → ∞ , λ c → B m .
That is, λ c = max ( B m , B ) .
D Proof of theorem
Proof.
Given that λ c c i,t ≥ c i − 1 ,t with some λ c ∈ R + , we can use c t,t λ t − i c as the upper bound on c
i,t with i = 1 , t , respectively.
Therefore,
[Equation]
where f ( λ ) =
c T,T T (cid:80) t =1 λ T − t is continuous on R + .
According to intermediate value theorem, there exists λ ∈ (0 , λ c ] such that c T,T T (cid:80) t =1 λ T − t =
T (cid:80) t =1 c t,T .
E Proof of theorem
Proof.
According to Theorem 1, there exists some λ i ∈ R + such that the summation of contribution stored between K i and K i +1 can be quantiﬁed as c K i +1 ,K i +
K i +1 (cid:80) t
=
K
i λ K
i +1 − t
i (after ignoring contributions before K i -th timestep for simplicity).
Let denote P ( λ )
=
K i +1 (cid:80) t
=
K
i λ K
i +1 − t , we have P (cid:48) ( λ ) > 0 , ∀ λ ∈ R + .
Therefore, P ( λ i ) ≥ P (cid:16) min i ( λ i ) (cid:17) .
Let C =
min
i ( c i,i ) and λ = min i ( λ i ) , the average contribution stored in a MANN has a lower bound quantiﬁed as I λ , where I λ =
C
[Equation]
F Proof of theorem
Proof.
The second-order derivative of f λ ( x ) reads:
[Equation]
We have f (cid:48)(cid:48) λ, ( x ) ≤ 0 with ∀ x ∈ R + and 1 > λ > 0 , so f λ ( x ) is a concave function.
Thus, we can apply Jensen inequality as follows:
[Equation]
[Equation]
Equality holds if and only if l 1 = l 2 = ...
=
l D +
=
T D +1 .
We refer to this as Uniform
Writing strategy.
By plugging the optimal values of l i , we can derive the maximized average contribution as follows:
[Equation]
When λ = 1 , I λ =
CT D +1 (cid:80) i
=1 l i = 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 (cid:48)(cid:48) λ
( x ) > 0 with ∀ x ∈ R + ; λ > 1 , f λ ( x ) is a convex function and Eq.
(12) ﬂips 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 +1 L , g λ ( L, D ) =
g λ ( x ) =
Cx ( λ 1 x − 1 λ − 1 ) .
We have g (cid:48) λ
( x ) =
Cλ (cid:16) 1 − λ 1 x (cid:17) ( x − ln λ ) > 0 with 0 < λ ≤ 1 , ∀ x ≥ 0 , so g λ ( T, D ) is an increasing function if we ﬁx 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.
G Summary of synthetic discrete task format
Task
Input
Output
Double x 1 x 2 ...x T x 1 x 2 ...
x T x 1 x 2 ...x T
Copy x 1 x 2 ...x T x 1 x 2 ...x T
Reverse x 1 x 2 ...x T x T x T − 1 ...x
Add x 1 x 2 ...x T x 1 + x T − 1x 2 + x T − 2... x (cid:98) T/ 2 (cid:99) + x (cid:100) T/ 2 (cid:101)
Max x 1 x 2 ...x T max ( x 1 , x
[Table]
H UW performance on bigger memory
[Table]
I Memory operating behaviors on synthetic tasks
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 reﬂected 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 10the 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 ﬁrst written location for the ﬁnal 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.
[Figure]
[Figure]
c)
[Figure] [Figure]
r
[Figure]
For the max task
[Figure]
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 suﬀers 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 ﬁrst and lasttimesteps 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 reﬂects 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 reﬂects 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,
[Figure]
y
[Figure]
+ A s i n ( 2 f x + )
[Figure]
i
[Figure]
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
[Table]
L Details on document classification datasets
[Table]
Yelp Review Polarity (Yelp P.)
[Table]
Yelp Review Full (Yelp F.)
[Table]
[Table]
Yahoo!
Answers (Yah. A.)
[Table]
[Figure]
r i t [Figure] g a t
[Figure]
M
Document classification detailed records
Model
AG
IMDb
Yelp P.
Yelp F.
UW 93.
Mean/Std 93.54 ± 0.
91.32 ± 0.
96.36 ± 0.
65.04 ± 0.
CUW 93.
Mean/Std 93.73 ± 0.
91.25 ± 0.
96.36 ± 0.
65.16 ± 0.