Title: SEQUENCE MODELLING WITH AUTO-ADDRESSING AND RECURRENT MEMORY INTEGRATING NETWORKS

Abstract: Processing sequential data with long term dependencies and learn complex transitions are two major challenges in many deep learning applications. In this paper, we introduce a novel architecture, the Auto-addressing and Recurrent Memory Integrating Network (ARMIN) to address these issues. The ARMIN explicitly stores previous hidden states and recurrently integrate useful past states into current time-step by an efficient memory addressing mechanism. Compared to existing memory networks, the ARMIN is more light-weight and inference-time efficient. Our network can be trained on small slices of long sequential data, and thus, can boost its training speed. Experiments on various tasks demonstrate the efficiency of the ARMIN architecture. Codes and models will be available.

1. INTRODUCTION:

Recurrent neural networks, such as the Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) and Gated Recurrent Unit (GRU) (Cho et al., 2014) have shown promising performance for processing sequential data. However, it's known that RNNs suffer from gradient vanishing problem. Moreover, as pointed out by Rae et al. (2016) , the number of parameters grows proportionally to the square of the size of the hidden units, which carry the historical information. Recent memorybased approaches exhibit potential to address these issues, by decoupling memory capacity from model parameters, and backpropagating the gradients through the memory.

Neural Turing Machine (NTM) (Graves et al., 2014) first emerged as a recurrent model that incorporates external memory abilities. NTM maintains a memory matrix, and at every time-step, the network reads and writes (with erasing) to the memory matrix using certain soft-attentional mechanism, controlled by an LSTM that produces read and write vectors. NTM and its successor, the Differentiable Neural Computer (Graves et al., 2016) , have shown success on some algorithmic tasks such as copying, priority sorting and some real-world tasks such as question answering. But one limitation of the NTM is that due to its smooth read and write mechanism, NTM has to do propagations on the entire memory, usually causing huge amount of memory consumption. To this end, Rae et al. (2016) proposes the Sparse Access Memory(SAM) network, by thresholding memory modifications to a sparse subset, i.e. all read and write operations are limited to several memory words. This allows memory and time efficient propagations while maintaining NTM's performance. However, these external memory models have relatively complicated memory addressing mechanisms, making them suitable for only large-scale memory. Moreover, the RNN in these models plays a simple role of being a controller, but the gradient vanishing problem of RNN itself is not given attention. In contrast to these networks, our network uses a light-weight read mechanism to address a small external memory, and relieves the gradient vanishing problem by directly concatenating memory word to output.

Inspired by prior memory models, efforts have been made to build a bridge between simple RNNs and complicated memory models. Kurach et al. (2015) propose the Neural Random-access Machines (NRAM) that can manipulate and dereference pointers to an external variable-size random-access memory. Danihelka et al. (2016) improve LSTM with ideas from Holographic Reduced Representations (Plate, 2003) that enables key-value storage of data. Grave et al. (2016) propose a method of augmenting LSTM by storing previous (hidden state, input word) pairs in memory and using the current hidden state as a query vector to recover historical input words. This method requires no backpropagation through memory and is well-suited to word-level language tasks. Grefenstette et al. (2015) ; Dyer et al. (2015) ; Joulin & Mikolov (2015) augment RNNs with a stack structure that works as a natural parsing tool, and use them to process algorithmic and nature language processing (NLP) tasks; nonetheless, the speed of stack-augmented RNNs is rather slow due to multiple push-pop operations at every time-step. Ke et al. (2018) propose the Sparse Attentive Backtracking (SAB) architecture, which recalls a few past hidden sates at every time-step and do "mental" backpropogations to the nearby hidden states with respect to the past hidden states. Gulcehre et al. (2017) propose the TARDIS network, which is an LSTM-resembled RNN that directly stores a fixed number of previous hidden states with learned key parameters for memory addressing. At every time-step, the network reads out one historical state r t from the memory, and uses it, along with the input x t and last hidden state h t-1 , to produce new hidden state h t , then overwrites h t to the location of r t . TARDIS optionally uses the gumbel-softmax estimator (Maddison et al., 2016; Jang et al., 2016) to sample the location of r t (we will explain this in section 2.1), so the whole network is differentiable. However, the TARDIS still involves some hand-crafted memory addressing method and its addressing mechanism requires a considerable amount of memory and time consumption.

Inspired by the TARDIS, we introduce the ARMIN architecture, a more light-weight and memorybased architecture with a very simple memory addressing mechanism and in-depth modification of the existing LSTM structure. Concretely, our contributions are as follows:

• We distill previous memory addressing methods and propose a simple yet effective memory addressing mechanism for the external memory, namely Auto-addressing, by encoding the information for memory addressing directly via the inputs x t and the hidden states h t-1 .

• We propose a novel recurrent cell that combines the gating advantages of LSTM and allows direct gradient backpropagation from output to memory. With only a single cell, it learns better representations than many hierarchical RNN structures.

• We show in our ablation study that the ARMIN is robust to small iteration lengths when training long sequential data, which enables training with large batch sizes and boost its training speed. We further demonstrate in char-level language modelling task that the ARMIN obtain 40% training speed gain than LSTM while keeping similar performance and memory consumption.

• We demonstrate competitive results on various tasks while keeping efficient time and memory consumption during training and inference time.

2. BACKGROUND:



2.1. GUMBEL-SOFTMAX ESTIMATOR:

Categorical distribution is a natural choice for representing discrete structure in the world; however, it's rarely used in neural networks due to its inability to backpropagate through samples (Jang et al., 2016) . To this end, Maddison et al. (2016) and Jang et al. (2016) propose a continuous relaxation of categorical distribution and the corresponding gumbel-softmax gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample. Specifically, given a probability distribution p = (π 1 , π 2 , ..., π k ) over k categories, the gumbel-softmax estimator produces an one-hot sample vector y with its i-th element calculated as follows:

EQUATION

where g 1 , ..., g k are i.i.d samples drawn from Gumbel distribution (Gumbel, 1954) :

EQUATION

)

and τ is the temperature parameter. In practice, we usually start at a high temperature and anneal to a small but non-zero temperature (Jang et al., 2016) .

3. AUTO-ADDRESSING AND RECURRENT MEMORY INTEGRATING NETWORK:

In this section, we describe the structure of the ARMIN network as shown in Figure 1 . It consists of a recurrent cell and a small external memory that stores historical hidden states. While processing from memory with an one-hot read vector produced via passing x t and h t-1 to a fully connected layer followed by a gumbel-softmax function. (b) The ARMIN cell receives x t , h t-1 and r t as inputs and outputs o t and h t . o t is passed to output layers, and h t is passed to next time-step. (c) h t is written to the previous location of r t . sequential data, the ARMIN performs reading from memory, cell processing and writing to memory operations in chronological order during each time-step. In the following subsections, we first explain the structure of the recurrent cell, and then discuss the read and write operations.

3.1. THE RECURRENT CELL OF ARMIN:

Inspired by classical LSTM structure (please refer to APPENDIX A for details in RNN and LSTM structures), we propose a novel memory-augmented recurrent cell structure, namely the ARMIN cell. At every time-step, it takes in an input x t , the last hidden state h t-1 and a recovered historical hidden state r t chosen by a read operation, and produces an output vector o t and the new hidden state h t . The computation process is as follows:

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

where

h t-1 , h t ∈ R d h , r t ∈ R dr , x t ∈ R di , and W ig ∈ R 2d h ×(di+d h +dr) , W go ∈ R (4d h +dr)×(di+d h +dr) , b ig ∈ R d h +dr , b go ∈ R 4d h +dr , o t ∈ R d h +dr .

We refer to d h as the hidden size of the recurrent cell of ARMIN. Usually we have d h = d r and allocate equal number of weight parameters for h t and r t .

In equation 3 ∼ 5, two gates are calculated to control the information flow for h t-1 and r t respectively, generating gated hidden state h g t-1 and historical state r g t ; using this method, we can filter out the irrelevant information for the current time-step. Then as shown in equation 6, we compute the input gate i t , forget gate f t , cell state g t and output gate o h t for the new hidden state just like in classical LSTM structure. Additionally, an output gate o r t for historical state r t is computed. Next in equation 7, we compute new hidden state h t that is the sum of h t-1 and cell state g t , leveraged by forget gate f t and input gate i t . Finally in equation 8, we calculate the output of this timestep, which is the concatenation of the gated contents from h t and r t . Using this method, we can selectively backpropagate gradients from the output to the historical state r t . We involve r t (or r g t ) in all gates' computations. Intuitively, the h t-1 acts as the old working memory, and r t is treated as the long-term memory. The cell processes them with the input x t to generate the new working memory h t and the output o t . More specifically, each r t is a summary of historical hidden states selected by auto-addressing mechanism. The ARMIN cell learns to recurrently integrate the summary of long-term information from r t into the working memory h t . We will demonstrate the efficiency of recurrent memory integration in language modelling tasks.

Unlike the LSTM that passes a tuple of (h t , c t ) to the next recurrence and outputs h t at the same time, the ARMIN only passes one hidden state h t and separately uses o t as output, which has 2 benefits: one is that it takes lower memory consumption to store one hidden state rather than two, and the other is we can decouple the information needed for the output o t , from the information (i.e. h t ) that is needed for memorizing and being used at later time-steps. Furthermore, we list 3 important differences between the ARMIN cell and TARDIS cell, please refer to APPENDIX B.1.

3.2. READ OPERATION WITH AUTO-ADDRESSING:

The ARMIN maintains a memory matrix M ∈ R nmem×d h , where the constant n mem denotes the number of memory slots. At each recurrence, the ARMIN chooses a historical state r t from memory according to the information in x t and h t-1 ,which is formulated as follows:

EQUATION

EQUATION

)

where

W s ∈ R nmem×(di+d h ) , b s ∈ R nmem ,

s t is a one-hot vector sampled by gumbel-softmax function, s t (i) denotes the i-th element of s t , M (i, :) denotes the i-th row of M .

As opposed to previous memory networks such as Santoro et al. (2016) and Gulcehre et al. (2017) , we don't use any extra usage vectors and learnable key parameters that manually encode memory addressing information to assist memory addressing. As our ablation study (Appendix E) shows, the hidden state h t-1 is sufficient to encode the historical memory accessing information. Furthermore, the TARDIS addressing mechanism requires concatenating x t and h t to every memory cell, which usually causes large extra memory and time consumption in practical scenes. Please refer to APPENDIX B.2 for more precise comparison regarding addressing mechanism of the ARMIN and TARDIS networks.

3.3. WRITE OPERATION:

After the reading and cell processing stages, the ARMIN writes the new hidden state h t to the memory M . Following the TARDIS network (Gulcehre et al., 2017) , we simply overwrite h t to the memory slot where we just read out the r t (for conditions where d h is not equal to d r , we first use a linear layer to transform h t from d h dimension to d r dimension). But at the initial time-steps, we write the hidden states to the empty memory slots, until all empty slots are filled with historical states. In this way, we can maximally preserve useful historical information, because the ARMIN cell can learn to copy the useful information from r t to h t , and then write h t to the previous location of r t . The read/write mechanism can be viewed as a form of a skip-connection of hidden states. By training the whole network, the ARMIN can easily learn long-term dependencies via direct access to historical hidden states. From this point of view, our network is similar to the Skip RNN proposed by Campos Camunez et al. ( 2018), which differs from our network in that they use a binary state update gate to select whether the state of the RNN will be updated or copied from the previous time-step.

4. EXPERIMENTS:

We evaluate our model on algorithmic tasks, character-level language modelling task and temporal action detection and proposal task. We compare our network with NTM (with an LSTM controller), TARDIS and vanilla LSTM networks. We also directly compare the auto-addressing against the TARDIS addressing mechanism, by replacing auto-addressing with the TARDIS addressing method in ARMIN. We refer to this network as ARMIN+TARDIS-addr in our experiments. Please refer to APPENDIX C for some universal setups used in all our experiments below.

4.1. ALGORITHMIC TASKS:

Along with the NTM, Graves et al. (2014) introduced a set of synthetic algorithm tasks which consists of copy, repeat copy, associative recall, N-gram and priority sort tasks. Here we use four out of five of these tasks (excluding the N-gram task) to examine if our network can choose correct time-steps from the past and effectively make use of them. Please refer to APPENDIX C.1 for more explanations and hyperparameter setups with respect to these tasks. Note that we have a strong LSTM baseline with about 4 times larger parameter count than memory networks, and the training of our copy and priority sort tasks are more challenging than the original tasks in Graves et al. (2014) , for having more than 2 times longer input and target sequences. The models are optimized to minimize the averaged binary cross-entropy loss on the target multi-bit binary sequences. Following Gulcehre et al. (2017); Campos Camunez et al. (2018) , in all tasks, we consider a task solved if the validation loss is at least two orders of magnitude below the initial cost which is around 0.70, i.e. the validation loss (the validation set is generated randomly) converges to less than 0.01 (with less than 30% in 10 consecutive validation of sharp losses that are higher than 0.01 ). In our evaluation, we are interested in if a model can successfully solve the task in 100k iterations and the elapsed training time and iterations till the model succeeds on the task. For all networks we ensure the similar total parameter counts. All important hyperparameters are shown in APPENDIX C.1. We run all experiments with batch size of 1 under 2.8 GHz Intel Core i7 CPU and report the wall clock time and the number of iterations for solved tasks of each model. For models that fail to converge to less than 0.01 loss in 100k iterations, we report the average loss of the final 10 validations (denoted in underlines) and the elapsed training time for 100k iterations. The results are shown in Table 1 . The training curves for each task are shown in APPENDIX D.

The results show that the ARMIN can solve 3 out of 4 tasks in a fast converge speed. By comparing the ARMIN with NTM, we observe that the NTM is able to almost solve all of the tasks, but with a much slower training speed compared to the ARMIN, for example, the training time of NTM is 5 times larger than ARMIN to solve the copy task. In fact, the training speed of ARMIN for each iteration is about 3 ∼ 4 times faster than the NTM's speed. We observe that the TARDIS and ARMIN+TARDIS-addr fail to solve 3 out 4 tasks in the given 100k iterations. By comparing them with the ARMIN, we confirm the efficiency of the auto-addressing of ARMIN, for bringing a fast We also find that the LSTM, when given enough parameter counts (4 times larger than memory networks), is able to solve the priority sort task quickly and converge to a better performance than TARDIS and ARMIN+TARDIS in other tasks. We believe there are 2 possible reasons for the poor performance of the TARDIS addressing mechanism: a) The tanh activation (see equation 21 in APPENDIX B.2) might cause the gradient vanishing problem when it cooperates with the gumbel-softmax function. b) The multiple inputs of x t , h t , M t [i], u t in equation 21 cause the network hard to initialize to keep equal input and output variances for the linear layer and might also cause gradient dispersion. The auto-addressing has neither of these two issues and can better integrately backpropogate gradients, therefore it can perform well and quickly on these tasks. However, we would like to point out 2 limitations of the auto-addressing: a)

The associative recall task shows the auto-addressing might have problems in establishing indirect references among memory data. b) The auto-addressing is relatively hard to generalize to unseen longer sequence in tasks such as copy and repeat copy due to its simplicity, whereas the NTM is able to do that. We leave these 2 issues for future work.

4.2. CHARACTER-LEVEL LANGUAGE MODELLING:

The character-level language modelling task consists of predicting the probability distribution of the next character given all the previous ones, and we use the widely used Penn Treebank and Hutter Prize Wikipedia (also known as enwik8) datasets for this task. We compare our network with vanilla LSTMs and other memory networks on Penn Treebank, and compare with some popular RNN variants on both datasets. We ensure similar parameter counts and hyperparameters for all memory networks. Following prior works, we apply truncated backpropagation through time(TBPTT) (Rumelhart et al., 1986; Elman, 1990) to approximate the gradients: at each iteration, the network predict the next 150 characters, and the hidden state h t and memory state M are passed to the next iteration. The gradients are truncated between different iterations. Please refer to APPENDIX C.2 for more experimental details.

The results are shown in table 2. On Penn Treebank dataset, our best performing network achieves competitive 1.202 BPC on Penn Treebank dataset, which is the best single cell performance that we are aware of, with relatively small parameter count. By comparing TARDIS and ARMIN+TARDIS-addr, we observe about 4 points of improvement, which shows the efficiency of the ARMIN cell. By comparing ARMIN+TARDIS-addr, we observe a further improvement of around 2 points, which shows the efficiency of the auto-addressing mechanism. We also find that the same architecture of NTM that performs well in algorithmic tasks fails to converge to a lower BPC than vanilla LSTM, even with various regularizations we add to NTM.

The results on Penn Treebank show our single ARMIN cell learns better representations than many hierarchical RNN structures, such as the HM-LSTM, 2-Layer HyperLSTM and 21 layer IndRNN.

Our network is outperformed by Recurrent Highway Hypernetwork and Fast-slow LSTM which are both multi-scale and deep transition RNNs and are state-of-the-art results without any auxiliary technics on this dataset. By saying "better representations", we refer to the concatenation of the gated contents from h t and r t as in equation 8. For example, if we remove the gated contents of r t from o t , the ARMIN undergoes a BPC performance drop from 1.202 to 1.220, which is still better than the best performing BPC of 1.24 of the LSTM. We believe the rest of the performance gain comes from the recurrent memory integration of the ARMIN cell, which also favors a paradigm of deep transition, as is shown by the success of the deep transition RNNs on this task. For more ablation study regarding the auto-addressing mechanism and the ARMIN cell please refer to APPENDIX E (where we show the robustness of ARMIN to small TBPTT length).

The result on enwik8 demonstrates a simple 2-layer ARMIN can achieve competitive BPC performance of 1.33, with less parameter count compared to the HyperLSTM and HM-LSTM. By constructing deeper ARMIN network or even combining with other multi-scale and hierarchical RNN architectures, we believe the performance can be further improved.

4.3. TEMPORAL ACTION DETECTION AND PROPOSALS ON THUMOS' 14:

In this subsection, we evaluate the ARMIN on a more complicated real-world task concerning video analysis, i.e. the temporal action detection and proposals, which consists of taking an input video and producing a set of temporal intervals that are likely to contain human actions (Buch et al., 2017) . We use the Single-Stream Temporal Action Proposal(SST) (Buch et al., 2017) as our framework, and evaluate the model on the THUMOS' 14 dataset (Jiang et al., 2014) , which contains 20+ hours of video with 200 train/validation and 213 test untrimmed video sequences. The SST consists of a GRU and a fully connected output layer. In our experiment, we replace the GRU encoder in SST with a single ARMIN cell, and we call the modified network MA-SST. For fairness, we also implement a modified SST with LSTM encoder and compare them with the original SST in Buch et al. (2017) . Please refer to APPENDIX C.3 for more experimental details. tIoUs. We believe the MA-SST has better performance than the SST mainly because of the direct access to historical states from which it can better locate the starting frame of an action.

5. TOWARDS MORE LIGHT-WEIGHT RECURRENT MEMORY NETWORKS:

It's known that the speed and computer memory consumption have important influences on the practicability of an RNN, especially when external memory is involved. We have shown in some of the algorithmic tasks that the ARMIN can be trained 3 4 times faster than the NTM in terms of wall clock time. Next we conduct a more thorough comparison experiment among the vanilla LSTM and memory networks under different hidden sizes. The experiment is based on the character-level language modelling task on Penn Treebank dataset, without the loss of generality. For a fair comparison, we implement all networks with near-optimal implementations using Pytorch, an efficient deep-learning framework. We use the aforementioned experiment setup, and keep all memory matrices the same size, i.e. 20 × d h . We run the experiment using single-precision floating point calculations under a Titan XP GPU that has 12GB memory space. The results are depicted in figure 3 . From the results we observe 2 important phenomena: a) The ARMIN consistently outperform other memory networks shown in the graph in terms of running speed both at training and inference stages, and the main contribution to this comes from the simple auto-addressing mechanism of the ARMIN. Moreover, at inference stage, we can replace the memory matrix with a list of discrete memory slots, and update memory by simply replacing the old hidden states with the new ones, furthermore, we can replace the gumbel-softmax function with argmax. Using these methods, ARMIN's inference speed obtains significant improvement than in training stage. b) The ARMIN outperforms TARDIS and ARMIN+TARDIS-addr in terms of training memory consumption, mainly because the TARDIS addressing mechanism requires concatenating x t and h t-1 to every memory cell for parallelization. At inference stage, we observe TARDIS have smaller memory consumption than ARMIN and ARMIN+TARDIS-addr, mainly because TARDIS has smaller parameter counts under the same hidden size. However, we would like to point out that under the same parameter counts the ARMIN usually outperforms TARDIS as we show in algorithm and language modelling tasks.

We also observe the NTM has the largest memory consumption at inference time, mainly because its complex addressing mechanism.

We have shown in the ablation study (see APPENDIX E) that our network is robust to small iteration length when trained with TBPTT. We realize that we can use this feature to enable large batch size, and boost ARMIN's training speed while keeping better performance than LSTM. To validate this benefit, we conduct an experiment to compare the performance and cost of ARMIN against our best performing LSTM. We also provide the test result of the same LSTM under TBPTT length T trunc = 100 for better comparing. The hyperparameter setups and results are shown in Table 4 . 

6. CONCLUSION:

In this paper, we have introduced the ARMIN, a light-weight and memory-augmented RNN architecture with a novel ARMIN cell. The ARMIN incorporates an efficient external memory with the light-weight auto-addressing mechanism. We demonstrate competitive performance of ARMIN in various tasks, and shows the generality of our model. Our ablation study suggests the efficacy of our external memory and addressing mechanism, and notably, our network is robust to shortlength TBPTT which enables using large batch size to speed up the training and further increase performance on sequence modelling tasks. Further research may lead to the efficient hierarchical and multi-scale structures of the ARMIN and its successful applications in Seq2Seq models and complicated reasoning tasks.

APPENDIX A RNN AND LONG SHORT-TERM MEMORY:

A recurrent neural network (RNN) is a class of neural network that recurrently processes a sequence of inputs {x 1 , x 2 , ..., x T }, and returns a sequence of outputs {y 1 , y 2 , ..., y T } where x t ∈ R di , y t ∈ R do , and d i , d o is the input size and output size, respectively. In a vanilla RNN, y i is given by following equations:

EQUATION

EQUATION

where the hidden state h t ∈ R d h is passed to the next recurrence. d h is the hidden size, and U , W , V , b, c are learnable parameters.

The LSTM (Hochreiter & Schmidhuber, 1997 ) is designed to overcome the gradient vanishing and exploding problem in vanilla RNN and learn long-term dependencies, by passing a tuple of (h t , c t ) to the next recurrence. Its computation process is defined as follows:

EQUATION

EQUATION

EQUATION

where 

• is the element-wise product, W 0 ∈ R 4d h ×(di+d h ) , b 0 ∈ R 4d h , i t , f t ,

APPENDIX B DIFFERENCES BETWEEN TARDIS AND ARMIN:

The major differences between ARMIN and TARDIS networks are the recurrent cell computation and read operation. We show these differences respectively in subsections below:

B.1 CELL COMPUTATIONS

The cell computation process of TARDIS is as follow:

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

where

h t-1 , h t , c t , i t , f t , o t , g t ∈ R d h , r t ∈ R dr , x t ∈ R di , w α , w β ∈ R di+d h +dr , and W ∈ R 3d h ×(di+d h +dr) , W g ∈ R d h ×(di+d h +dr) , b ∈ R 3d h , b g ∈ R d h .

The output of TARDIS is the concatenation of h t and r t .

There are 3 important points that distinguish the ARMIN cell from TARDIS cell, which are as follows:

• The TARDIS uses two binary scalar gates α t , β t to control the information flow in equation 17.

The two scalar gates are activated using gumbel-sigmoid (which is similar to gumbel-softmax and the output is close to binary), whereas the ARMIN uses two soft vector gates to control the information flow, which is more flexible and enable larger model capacity. For example, in the charlevel language modelling experiments, if we replace the sigmoid activations in equation 3 with the element-wise gumbel-sigmoid, the BPC performance immediately drops from 1.202 to 1.235.

• The TADIS apply information control only when calculating the cell state g t , whereas the ARMIN apply information control immediately after the read operation. Our main motivation to this change is that if we don't filter out irrelevant information in r t and h t at first, these noisy information will cause the weight parameters hard to learn. Furthermore, if the read operation chooses the wrong past state, we can block the information flow for r t at the first place.

• The ARMIN has an extra gate o r t to control the output from r t as in equation 8, which the TARDIS doesn't have. If we remove this gate and directly output r t like in TARDIS, we get a BPC performance drop from 1.202 to 1.219 in char-level language modelling task.

Finally, despite that the ARMIN has 3 more vector gates than the TARDIS, we have shown in the pMNIST and Penn Treebank experiments that the ARMIN can still outperform the TARDIS under similar parameter counts.

B.2 READ OPERATION:

The memory matrix M t of TARDIS has disjoint address section

A t ∈ R nmem×d k and content section C t ∈ R nmem×dr , M t = [A t ; C t ].

The controller reads both the address and the content parts of the memory, but it will only write into the content section of the memory.

The TARDIS paper explored two ways to sample the read locations--using reinforce and using gumbel-softmax function, respectively. The gumbel-softmax way shows a better performance and its read operation can be formulated as follow:

EQUATION

EQUATION

)

where {a, W γ , b γ } are learnable parameters (note that in our experiments, we always set the dimension of a to about 1/4 of the hidden size) and u t is a hand-crafted usage vector which denotes the frequency of accesses to each cell in the memory.

There are two important differences with respect to TARDIS addressing mechanism and our autoaddressing mechanism:

• In practice, equation 21 requires concatenating x t and h t to every memory cells M t [i] for parallelization, which often causes significant extra memory and time consumption compared to the auto-addressing of ARMIN, as we show in the algorithmic tasks and char-level language modelling task.

• The auto-addressing omits the learnable address section A t , the hand-crafted u t vector, the tanh activation and an auxiliary vector a compared to TARDIS addressing mechanism. The TARDIS also requires manually subtracting 100 in unnormalized probability at last read location to avoid repeated read, whereas the auto-addressing doesn't need to. With these simplifications, we found the auto-addressing mechanism easier to train and it leads to a much faster converge speed, as we show in the algorithmic tasks. We also observe a BPC performance growth from 1.223 to 1.202 in char-level language modelling task when we switch from ARMIN+TARDIS-addr to ARMIN with auto-addressing.

APPENDIX C EXPERIMENT SETUPS:

We first describe some universal setups in all our experiments (unless we mention them again with different setup). Our model can be trained end-to-end, due to the utilization of gumbel-softmax function. As pointed out in Maddison et al. (2016) , when the temperature τ ≤ (n -1) -1 , the density function of gumbel-softmax estimator is log-convex, which gives us a good guarantee for optimization. However, smaller τ will also bring higher variance of the gradients (Jang et al., 2016) . So in all experiments, we initialize the reciprocal of τ as 1, and increase it by 1 after each epoch, until the reciprocal equals to n -1. This scheme has shown better performance than other schemes.

We apply Layer Normalization (Ba et al., 2016) to all networks. In the ARMIN case, layer norm is applied to equation 3, 6 and 7, which we find important to regulate the hidden state when scale grows. To avoid overfitting, we apply dropout (Srivastava et al., 2014) to the input and output layers, and Zoneout (Krueger et al., 2016) is applied for recurrent connections. The networks are trained with Adam optimizer (Kingma & Ba, 2014) . All weight matrices in RNNs are initialized to orthogonal matrices, and the bias of the forget gate f t is initialized to 1. For all models using TARDIS addressing mechanism, we set the memory key size to about 1/5 of the memory content size.

C.1 ALGORITHMIC TASKS:

In all algorithmic experiments below, following Graves et al. (2014) , we apply RMSProp optimizer (Tieleman & Hinton, 2012) with a momentum of 0.9. The learning rate is 0.0001 and gradient norms are clipped to 10.0. No regularization is used as the size of the networks is small. Following Graves et al. (2014) , we have learned bias of initial hidden states and memory matrices for all memory networks. For TARDIS, ARMIN+TARDIS-addr and ARMIN we increase the reciprocal of temperature τ by 1 for every 200 iterations. The hyperparamters for models and tasks are shown in Table 5 and Table 6 , respectively. We explain these tasks in detail in subsections below:

C.1.1 COPY

The copy task tests whether a recurrent model can correctly store and recall a long sequence of arbitrary information. In our experiment, the networks are given randomly generated eight-bit binary sequences with its length ranging from 1 to 50, followed by a finish indicator. Then the networks are asked to output the same sequence with the input.

C.1.2 REPEAT COPY:

The repeat copy task is a repeat version of the copy task. The main motivation was to see if a recurrent model could memorize input information and repeatedly reuse them at later time-steps. In our experiment, the networks are asked to repeatedly copy the randomly generated input sequence with length ranging from 1 to 10 to output for 1 to 10 times.

C.1.3 ASSOCIATIVE RECALL:

This task test whether a memory network has the ability to retrieve memory data indirectly. The networks are given a sequence of items and then a query with one of the items, and are asked to output the subsequent item. In our experiments, each item consists of three six-bit binary vectors, and the number of items are ranging from 2 to 6. We use the train/validation/test split outlined in Mikolov et al. (2012) . The train/validation/test batch size are 128/64/1 respectively. Following prior works on this task (Ha et al., 2016; Suarez, 2017; Mujika et al., 2017) , the model is optimized over the cross-entropy loss between the predictions and the training labels, and evaluated using a bits-per-character measure. We use a single ARMIN cell of different hidden sizes, with a fully connected layer to output the predictive probabilities. The input embedding size d i is set to 128. The number of memory slots n mem is 20. The zoneout probability is 0.3. We train the networks for 200 epochs with a learning rate of 0.002, and decay with a factor of 10 at the last 20 epochs. We clip the gradients to a maximum norm of 1.0.

C.2.2 HUTTER PRIZE WIKIPEDIA DATASET:

We use the train/validation/test split outlined in Chung et al. (2015) . The train/validation/test batch size are 128/64/1 respectively, and the vocabulary size is 205. The zoneout and dropout probabilities are 0.3 and 0.2, respectively. The input embedding size is 256. Our ARMIN network has 2 layer and 800 units at each layer, connected by a linear layer to convert 1600-dimensional features from layer 1 to 256 dimension as input to layer 2. n mem for each layer is set to 10. We train the networks for 50 epochs with a learning rate of 0.001, and decay with a factor of 10 at the last 10 peochs. We clip the gradients to a maximum norm of 1.0.

C.3 TEMPORAL ACTION DETECTION AND PROPOSALS ON THUMOS' 14:

The SST takes in sequential inputs of a video and outputs the existing probability predictions of action proposals with different time lengths that end at each time-step, and uses a multi-label loss to optimize its parameters. The sequential inputs are 4096-dimensional features extracted from every 16 consecutive video frames using C3D network (Tran et al., 2015) , followed by PCA to downsize the input feature dimensionality.

Our experiment setup is basically the same as the SST paper, except that we use a simple fully connected layer to downsize the input features instead of using PCA. The input size of the two networks are both set to 256. The MA-SST have 256 hidden units, while the SST with LSTM have 512 hidden units. This ensures the 2 networks have roughly the same parameter count. The dropout and zoneout probabilities are both set to 0.2 and 0.3. n mem of the MA-SST is set to 35. We train the networks for 300 epochs with initial learning rate of 0.002, and decay with a factor of 10 at the last 60 epochs. • Random Read The ARMIN randomly chooses a historical state while reading. The result shows a BPC of 1.224, compared to the 1.202 BPC of normal ARMIN. It proves that the read regime we proposed in section 3.2 is effective and can choose more useful historical state than random read.

• Queue-style Write While writing, the ARMIN follows a fashion of first-in-first-out, i.e. always chooses the oldest memory slot to write. It makes the memory M stubbornly store the hidden states of the last 20 time-steps. The result shows a BPC of 1.352, which proves our normal read/write mechanism has much more flexibility and can learn long-term patterns.

• Independent Read/write While writing, the ARMIN doesn't write to the last read memory slot, but independently choose a memory slot to overwrite according to x t and h t , in a similar fashion with the read operation we presented in section 3.2. The result shows a BPC of 1.245, proving that the normal writing regime in the ARMIN can better learn long-term patterns and maximally preserve historical information, as we explain in section 3.3.

• Smooth Read/write We replace the gumbel-softmax function with the vanilla softmax function, so the ARMIN reads and writes in a smooth way. The results show a minor performance drop of 1.212 BPC. The result shows that our discrete read/write regime is still better. Moreover, the smooth read/write regime would consume much more memory at both training and inference stages.

:

As implemented inHa et al. (2016).

Our implementation. We don't see performance growth when we further increase the hidden size.

tIoU denotes temporal Intersection over Union between a proposal and its maximally overlapped ground truth action. Applying bigger tIoU threshold make the generated proposals higher in quality but less in quantity.

Figures and Tables:

Figure1: The ARMIN structure. At each time-step, the ARMIN performs read operation, cell processing and write operation in chronological order: (a) It reads out a historical hidden state r t from memory with an one-hot read vector produced via passing x t and h t-1 to a fully connected layer followed by a gumbel-softmax function. (b) The ARMIN cell receives x t , h t-1 and r t as inputs and outputs o t and h t . o t is passed to output layers, and h t is passed to next time-step. (c) h t is written to the previous location of r t .

Figure 3: The running speed and memory consumption at the training and inference stages (n mem = 20). The solid blocks shows the memory consumption and the curves denote the running speed in characters/s. (a) shows the training stage, and (b) shows the inference stage (note that ARMIN+TARDIS-addr has basically the same memory consumption with ARMIN, so it's not shown in the graph).

Figure 4: Copy task.

The average elapsed time(min) and iterations(k) of different networks till task solved in given 100k iterations. The wall clock training time for 100k iterations and the average loss of final 10 validations (denoted in underlines) are shown for the unsolved tasks.

Table content: <table><tr><td/><td colspan="2">Copy</td><td colspan="2">Repeat Copy</td><td colspan="2">Associative Recall</td><td colspan="2">Priority Sort</td></tr><tr><td/><td>time</td><td>iter.</td><td>time</td><td>iter.</td><td>time</td><td>iter.</td><td>time</td><td>iter.</td></tr><tr><td>LSTM</td><td>66</td><td>0.359</td><td>45</td><td>0.016</td><td>34</td><td>0.325</td><td>32</td><td>37.0</td></tr><tr><td>NTM</td><td>32</td><td>12.4</td><td colspan="2">171 0.014</td><td>22</td><td>19.6</td><td colspan="2">360 0.012</td></tr><tr><td>TARDIS</td><td colspan="2">157 0.451</td><td>97</td><td>0.166</td><td>66</td><td colspan="2">0.330 133</td><td>62.6</td></tr><tr><td colspan="3">ARMIN+TARDIS-addr 152 0.410</td><td>95</td><td>0.297</td><td>64</td><td colspan="2">0.336 151</td><td>71.4</td></tr><tr><td>ARMIN</td><td>6</td><td>7.6</td><td>36</td><td>67.8</td><td>33</td><td>0.052</td><td>37</td><td>33.4</td></tr></table>

Bits-per-character on Penn Treebank and enwik8 test set.

Table content: <table><tr><td>PTB</td><td>enwik8</td></tr></table>



Table content: <table><tr><td colspan="3">: Comparison of proposal generation per-</td></tr><tr><td colspan="3">formance in terms of recall at 1000 proposals.</td></tr><tr><td colspan="3">We choose tIoU=0.6 and 0.8 to compare for con-</td></tr><tr><td colspan="3">sistency with Buch et al. (2017). Our MA-SST</td></tr><tr><td colspan="2">achieves highest recall at tIoU=0.8.</td><td/></tr><tr><td>Model</td><td colspan="2">tIoU=0.6 tIoU=0.8</td></tr><tr><td>SST (Paper)</td><td>0.920</td><td>0.672</td></tr><tr><td colspan="2">SST (We implement) 0.897</td><td>0.732</td></tr><tr><td>MA-SST(Ours)</td><td>0.911</td><td>0.759</td></tr></table>

Comparison for performance and cost of different setups on Penn Treebank dataset.Figure 2: Comparison of proposal generation performance. For data generation, we use the code offered in Buch et al. (2017) for consistency. (Left) The average number of proposals v.s. average recall for tIoU ≥ 0.5. The MA-SST outperforms our SST when average number of proposals is less than 500, and outperforms the original SST when average number of proposals is more than 500. (Right) Different tIoU v.s. recall of 1000 average proposals. The MA-SST strongly outperforms the original SST and outperform our SST by a maximal margin of 2.7%.

Table content: <table><tr><td>Model</td><td colspan="2">ARMIN</td><td colspan="2">LSTM</td></tr><tr><td>Setup</td><td>1</td><td>2</td><td>1</td><td>2</td></tr><tr><td>Hidden size</td><td>500</td><td>550</td><td>1k</td><td>1k</td></tr><tr><td>n param (M)</td><td>4.02</td><td colspan="3">4.81 4.79 4.79</td></tr><tr><td>n mem</td><td>5</td><td>10</td><td>-</td><td>-</td></tr><tr><td>T trunc</td><td>50</td><td>50</td><td colspan="2">100 150</td></tr><tr><td>batch size</td><td>384</td><td>300</td><td colspan="2">128 128</td></tr><tr><td colspan="2">Memory(GB) 3.49</td><td colspan="3">3.56 2.36 3.27</td></tr><tr><td>Speed (chars/s)</td><td>98k</td><td>75k</td><td colspan="2">71k 70k</td></tr><tr><td>BPC</td><td colspan="4">1.238 1.226 1.27 1.24</td></tr></table>

By comparing ARMIN setup 1 with LSTM setup 2, we observe that with only 6.7% more memory consumption, we obtain 40% training speed gain while keeping a slightly better BPC performance and 16% less parameter count; by comparing ARMIN setup 2 with LSTM setup 2, we observe that with only 8.8% more memory consumption, we obtain about 1.5 points of BPC performance gain and 7.1% training speed gain under similar parameter count.

Table content: <table><tr><td/><td>90000</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>11</td></tr><tr><td/><td colspan="2">Speed(NTM)</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>Running Speed(Chars/s)</td><td colspan="4">Speed(TARDIS) 22500 45000 67500 Speed(ARMIN+TARDIS-addr) Speed(ARMIN) Speed(LSTM) Memory(LSTM) Memory(ARMIN) Memory(NTM) Memory(TARDIS) Memory(ARMIN+TARDIS-addr)</td><td/><td/><td/><td/><td/><td/><td>2.75 5.5 8.25</td><td>Memory Consumption(GB)</td></tr><tr><td/><td>0</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>0</td></tr><tr><td/><td>100</td><td>200</td><td>300</td><td>400</td><td>500</td><td>600</td><td>700</td><td>800</td><td>900</td><td>1000</td><td>1100</td></tr><tr><td/><td/><td/><td/><td/><td colspan="3">Hidden Size</td><td/><td/><td/></tr></table>

g t o t are the input gate, forget gate, cell state, output gate at time t, respectively, and they control the information flow in the LSTM.

Table content: <table/>

Hyperparameters of the models for algorithmic tasks.

Table content: <table><tr><td>Model</td><td>Hidden Size</td><td>memory Size</td><td>Param Count</td></tr><tr><td>LSTM</td><td>300</td><td>-</td><td>376k</td></tr><tr><td>NTM</td><td>120</td><td>128×20</td><td>88k</td></tr><tr><td>TARDIS</td><td>120</td><td>50×32</td><td>90k</td></tr><tr><td>ARMIN+TARDIS</td><td>100</td><td>50×32</td><td>91k</td></tr><tr><td>ARMIN</td><td>100</td><td>50×32</td><td>90k</td></tr></table>

Hyperparameters of the algorithmic tasks.

Table content: <table><tr><td>Task</td><td>Setup</td></tr><tr><td>Copy</td><td>copy length = 1∼50</td></tr><tr><td>Repeat Copy</td><td>copy length = 1∼10 number of repeats = 1∼10</td></tr><tr><td>Associative Recall</td><td>item size = 3×6 item number = 2∼6</td></tr><tr><td>Priority Sort</td><td>Sort length = 40</td></tr></table>

C.1.4 PRIORITY SORTThis task test whether a recurrent model can do priority sort--a classic elementary algorithm. This task is more challenging than previous ones as it not only require memorizing the input sequences but also making efficient use of the mutual relations of the items in the input sequences. In our experiments, the networks are given 40 eight-bit random binary keys with scalar priority values in a sequence, and are asked to output the top 30 keys in descending order of priority.

Table content: <table><tr><td>C.2 CHARACTER-LEVEL LANGUAGE MODELLING ON PENN TREENBANK DATASET</td></tr><tr><td>C.2.1 PENN TREEBANK DATASET</td></tr></table>

Peer review (iclr2019):
 This paper introduces a memory-augmented RNN (MARNN) which aims at being lightweight and differentiable. In a nutshell, authors propose to augment a LSTM-type architecture with several memory cells. At each time-step, MARNN retrieves one memory cell, updates his state, and updates the memory cell content. To learn the retrieval operation that requires discrete addressing, authors rely on the Gumbel-Softmax. Authors evaluate their approach on PennTreeBank character level modelling where they demonstrate competitive performances. They also report state-of-art performance on the Thumos dataset. The paper is overall clear and pleasant to read. Authors highlight that MARNN is more lightweight compared to existing memory networks. MARNN can indeed retrieves only memory cells at inference. However, since MARNN uses a Gumbel-Softmax to train the discrete addressing scheme, it is it not clear if there is any advantage in term of memory and computation of MARNN relatively to other network during training? It would be nice to compare the computation time/memory usage of MARNN with other memory augmented network such as TARDIS, NTM or Memory Network during training and inference. Another claim is that MARNN can possibly boost training speed by reducing the lengths of TBTT. But MARNN also haves a training time overhead as showed in Figure 2. How does the overall training time/performances of MARNN with TBPTT of 50 compared to a LSTM with TBPTT of 100/150? The writing can be sometime a bit imprecise. For instance authors say that MARNN “learns better representations that many hierarchical RNN structure”. I agree that MARNN outperforms in term of accuracy, however, it is not clear what the author are referring to by “better representation” of the MARNN hidden state? Performance gain of MARNN could also be due to the external memory which allows to retain more information of the input? In addition, it would be nice to precise which type of hierarchical RNN structure MARNN does (or doesn’t) outperform. Another claim is that MARNN can “easily learn long-term dependencies”. While this is reasonable, I am unsure that the empirical evaluation support this. It would be nice to show how the gradients backpropagated through time behave in practice to support this claim? Memory-augmented network are a very important research directions and the MARNN architecture is interesting. However, it is not entirely clear to me what is the main advantage of MARNN relatively to other memory networks network such as TARDIS, NTM or Memory Network for training and/or inference. Although authors do compare with TARDIS, further comparison with the other networks and in term of computation time and memory could help clarify those points.