M EMORY A RCHITECTURES IN R ECURRENT N EURAL N ETWORK L ANGUAGE M ODELS
A BSTRACT
We compare and analyze sequential, random access, and stack memory architectures for recurrent neural network language models.
Our experiments on the Penn
Treebank and Wikitext-2 datasets show that stack-based memory architectures consistently achieve the best performance in terms of held out perplexity.
We also propose a generalization to existing continuous stack models (Joulin & Mikolov, 2015; Grefenstette et al., 2015) to allow a variable number of pop operations more naturally that further improves performance.
We further evaluate these language models in terms of their ability to capture non-local syntactic dependencies on a subject-verb agreement dataset (Linzen et al., 2016) and establish new state of the art results using memory augmented language models.
Our results demonstrate the value of stack-structured memory for explaining the distribution of words in natural language, in line with linguistic theories claiming a context-free backbone for natural language.
I NTRODUCTION
Sequential recurrent neural networks such as LSTMs (Hochreiter & Schmidhuber, 1997) are the basis of state-of-the-art models of natural language in various tasks.
They effectively learn to capture dependencies between events separated in time by learning to store and retrieve information in a hidden state.
However, the ability of these methods to discover long-term dependencies is limited by the capacity of their hidden state and the difﬁculty of propagating reliable gradients.
For example,
LSTM language models have been shown to struggle to capture non-sequential syntactic dependencies in complex sentences without explicit supervision (Linzen et al., 2016).
As an illustration of the kind of dependencies they have difﬁculty learning, in the sentence, the loss of basic needs providers emigrating from impoverished countries has a damaging effect , correctly predicting singular has rather than its plural form have requires that the LSTM have learned that it depends on the subject, in this case the ﬁrst noun ( loss ) rather than any of the intervening non-subject nouns, such as countries .
Linzen et al.
(2016) show that LSTM language models fail to capture this kind of dependencies, especially as the number of attractors (underlined) increases.
Attempts to improve language models’ ability to capture non-local dependencies have recently been undertaken by introducing an external memory components.
These include (i) a soft attention mechanism (Daniluk et al., 2017) and (ii) an explicit memory block or cache model (Tran et al., 2016;
Grave et al., 2017).
However, since very local context is often most highly informative for predicting the next word, existing memory-augmented RNN LMs use memory just to store information about local context (Daniluk et al., 2017).
In this work, we compare several memory architectures for recurrent neural network language models.
Since our goal is to evaluate how well these types of memory architectures learn long term and syntactic dependencies, we focus on language models that are static as opposed to non-static models such as neural cache (Grave et al., 2017) and dynamic evaluation (Krause et al., 2017) that can update their distribution at test time .
We consider increasing the capacity of a purely sequential memory model by increasing the capacity of an LSTM, a random access memory model as typiﬁed by an attention-based LSTM, and a new variant of a stack augmented recurrent neural network.
Unlike random access memory models, a stack has a built-in bias to discover hierarchical structures that are important in language.
A continuous stack memory has been proposed to improve recurrent neural networks (Joulin & Mikolov, 2015; Grefenstette et al., 2015), although it has never been carefully evaluated in benchmark language modeling experiments.
In the only set of results for a stack augmented recurrent language model, Joulin & Mikolov (2015) show that a stack augmented vanilla RNN outperforms a standard RNN and is comparable to an LSTM.
We augment an LSTM with a stack memory and perform thorough comparisons to evaluate its efﬁcacy as a language model.
In contrast to prior work, our continuous stack allows for push, stay, and a variable number of pop operations at each time step (multiple pop operations are useful in modeling natural language sentences since while only a single new word is presented at each time step, multiple syntactic units may come to an end concurrently).
Our motivating hypothesis is that allowing the memory to dynamically store and retrieve contextual information with a stack will drive the model to use the memory to learn dependencies that are difﬁcult to capture by a sequential model.
Sequential memory has an easier time learning local dependencies, but it often fails to capture long term dependencies.
Random access memory models capture longer range dependencies (i.e., proportional to the window size), but the learner has to infer these from data without any informative structural bias.
We hypothesize that introducing a more appropriate inductive bias will make it easier for the model to learn long range and structurally meaningful dependencies, given that the variance in learning such dependencies can be high.
Linguistic insights reveal that one possible inductive bias is in the form of a hierarchical nested structure that captures syntactic dependencies.
Stack memory models provide a natural way for capturing hierarchical structures, providing an easier path for gradients to ﬂow to particular locations in the past.
Our main contributions in this paper are as follows:
• We thoroughly evaluate the efﬁcacy of a stack augmented RNN as a language model and propose a more expressive extension to existing stack models (§2.1).
• We compare how a recurrent neural network uses a stack memory, a sequential memory cell (i.e., an LSTM memory cell), and a random access memory (i e., an attention mechanism) for language modeling.
Experiments on the Penn Treebank and Wikitext-2 datasets (§3.2) show that both the stack model and the attention-based model outperform the LSTM model with a comparable (or even larger) number of parameters, and that the stack model eliminates the need to tune window size to achieve the best perplexity.
• We assess the ability of these memory models to discover long range structural dependencies commonly encountered in natural language using the subject-verb agreement dataset (Linzen et al., 2016).
We achieve new state of the art results and show that the gap in accuracy between a sequential or random access memory model with a stack model gets bigger as the dependencies become more complex (i.e., number of attractors increases; §3.3).
We also analyze the stack and ﬁnd that the model tends to use it to enhance its sequential memory component in high entropy prediction contexts (§3.4).
M ODEL
We consider a language modeling problem where the goal is to predict the next word x t given previously seen context words x 0 , . . .
, x t − 1 .
We represent each input word x by its D -dimensional embedding vector x ∈ R D .
Our base model is an LSTM that computes a hidden state at timestep t as follows:
[Equation]
Sequential memory.
An LSTM has a sequential memory cell c to store and retrieve information that is regulated by the input, output, and forget gates.
In order for long-term contextual information to be used in the future, it has to pass through these gates for multiple timesteps.
Random access memory.
One common approach to retrieve information from the distant past more reliably is to augment the model with a random access memory block via an attention based [Figure]
K
K
[Figure] x t [Figure] x t +
[Figure]
method.
In this model, we consider the previous K states as the memory block, and construct a memory vector m t by a weighted combination of these states: m t = t − 1 (cid:88)
i =
t − K a i h
i , where a i ∝ exp( w m,i h i + w m,h h t )
Such method can be improved further by partitioning h into a key, value, and predict subvectors (Daniluk et al., 2017).
Given the LSTM hidden state h t and the memory state m t , we combine them using a simple function:
[Equation]
to get the ﬁnal representation ˜h t .
We compute the probability of predicting the next word as p ( x t | x <t )
∝ exp( x (cid:62) t ˜h t + b y,x t ) , where we follow Inan et al.
and reuse the word embedding matrix X as the softmax parameters.
Stack memory.
In this work, we propose to augment a recurrent LSTM language model with a stack memory M that has three basic operations:
• P USH :
Push the current hidden state h t onto the stack.
• P OP :
Remove the top element of the stack.
• S TAY :
Keep the stack unchanged.
Figure 1 shows an illustration of our stack augmented RNN.
We describe the stack memory in details in the followings.
2. M ULTIPOP A DAPTIVE C OMPUTATION S TACK
Our stack is a multipop adaptive computation stack—it learns how many POP operations need to be performed before predicting an output.
In previous work (Joulin & Mikolov, 2015; Grefenstette et al., 2015), at every timestep t , the job of the memory (stack) controller is to decide whether (i) to push the current state (either h t or x t ) onto the stack, (ii) to pop the top element of the stack m 0 , or (iii) to stay and keep the stack state unchanged.
The stack of Joulin & Mikolov (2015) is primarily designed as a single computation stack that performs one of the available operations at every timestep.
In order to capture long-term dependencies, the stack learns to carry the information across multiple timesteps by mainly relying on the LSTM hidden states for predictions in between and keeping the state of the stack the same (i.e., by choosing to stay), or by pushing and popping the same number of times in between these timesteps.
While this promotes discoveries of hierarchical dependencies, the kind of hierarchical dependencies that it can discover is limited.
A multipop stack, on the other hand, has greater ﬂexibility since there are more ways to manipulate its state at each timestep.
The stack of Grefenstette et al.
(2015) implicitly allows multiple pop operations in a single timestep by setting the pop weights to be greater than one.
However, the controller makes this decision based only on the element at the top of the stack (along with the input and the current hidden state), making it less plausible to know whether more than one pop operations are needed since it does not look at other elements of the stack.
Our formulation of the multipop operations is more intuitive and takes inspirations from adaptive computation time (Graves, 2017).
Concretely, consider a stack memory with K elements.
In all our experiments, we limit the size of the stack to K = 10 for computational considerations.
If the stack requires more than K elements, the bottom element of the stack is removed to make space for the new element, which is added on the top of the stack.
In a single computation stack, a feedforward policy network is used to compute the probability of choosing an action a ∈ { STAY , PUSH , POP } .
In our stack, we also use a feedforward policy network, but the number of possible POP s is k ∈ { 0 , 1 , . . .
, K } .
Denote the current top two elements of the stack after performing k pops by m k,m k, 1 , and the state of the stack after k pops by STAY k
(i.e., do k pops and stay) or PUSH k
(i.e., do k pops and push the current hidden state h t ).
We compute the probability of choosing an action recursively:
[Equation]
As a base case, we have p ( POP − 1 | x t , m k − 1 , 0 , m k − 1 , 1 )
= 1 .
To ensure that the probability sums to one, we set p ( POP
K [Equation] | x [Equation]
[Equation] , 0 , m
[Equation]
= 0 . The ﬁnal stack state is then computed as: [Equation] where M ∗ k is the stack state after performing k POP and PUSH or STAY .
Denote the top of the ﬁnal stack at timestep t as m t .
The ﬁnal representation is
[Equation]
We propose to treat this stack as a fully differentiable continuous stack.
Alternatively, our stack can also be treated as a discrete stack and trained with reinforcement learning (e.g., with REINFORCE ;
Williams, 1992).
In this case, instead of summing over all possible stack states, we sample according to the probabilities.
However, such methods can have slow convergence due to high variance.
We include comparisons to discrete and continuous single computation stacks in our experiments (§3).
(Graves, 2017) consider the number of computations as “thinking time”, where they show that their models use more computation time for more difﬁcult predictions.
In our work, the number of computations is related to how further back we need to look back when making a prediction at a given timestep.
Note that when we decide to push, we push the current hidden state h t onto the stack.
Since this operation is performed before making a prediction at every timestep, it is possible to use the stack to increase the number of parameters for some predictions (i.e., by pushing h t and immediately use it to compute ˜h t
= W h,h h t + W h,m m t , because immediately after a push m t = h t ).
As a result, our stack is also related to variable computation recurrent networks Jernite et al.
(2017) that decide the number of dimensions to be used at each timestep.
E XPERIMENTS 3. S ETUP
We compare the following methods in our experiments:
• Sequential memory: 650-dimension and 920-dimension vanilla LSTMs (650 or 920 for both the word embedding and the LSTM hidden size).
• Random access memory: a 650-dimension attention-based LSTM with attention size K = { 1 , 3 , 5 , 10 , 15 } .
• Stack memory: a single computation discrete or continuous stack, or a multipop adaptive computation continuous stack on top of a 650-dimension LSTM.
Following Inan et al.
, we tie word embedding and word classiﬁer layers and apply dropout to these layers with probability 0.6 (value chosen based on preliminary experiment results).
We also use recurrent dropout (Semeniuta et al., 2016) and set it to 0.1.
We perform non-episodic training with batch size 32 using RMSprop (Hinton, 2012) as our optimization method.
We tune the RMSprop learning rate and (cid:96) 2 regularization parameter for all models on a development set by random search from [0 . 004 , 0 . 009] and [0 . 0001 , 0 . 0005] respectively, and use perplexity on the development set to choose the best model.
3. P ERPLEXITY
We use standard language modeling datasets, the Penn TreeBank (PTB) and Wikitext-2 (Wik-2) corpora to evaluate perplexity.
Our main results are summarized in Table 1, where we also show comparisons with previous work on these datasets.
Our basic LSTM model is comparable to some of the best LSTM models.
The results show that increasing the sequential memory capacity by increasing the hidden size improves performance.
However, the improvement is not as signiﬁcant as adding random access or stack memory.
The best attention model is the one with K = 10 and K = 15 on PTB and Wik-2 respectively, highlighting the necessity to tune to get the optimal window size.
Our results generally agree with Daniluk et al.
(2017) that show that increasing the attention size generally improves performance up to a certain threshold.
While both the discrete and continuous single computation stack models perform reasonably well, the discrete model underperforms the continuous model on Wik-2.
Recall that we use REINFORCE to learn the optimal discrete stack operations.
We leave it to future work to investigate whether better techniques can be used to improve the performance of the discrete stack.
The best model on both datasets is consistently the multipop stack.
Overall perplexity on these datasets is strongly dominated by words that require little to no long term dependencies, making it difﬁcult to assess when memory helps.
In the next section, we look into a speciﬁcally designed linguistic task to get a better understanding of these memory models.
3. S YNTACTIC D EPENDENCIES
We evaluate these memory models for learning syntax-sensitive dependencies on the number prediction dataset from Linzen et al.
In this dataset, the model is given a sentence up to—but not including—its verb, and the goal is to predict the number of the following verb (singular or plural).
For example, given a sentence preﬁx with different numbers of intervening nouns:
• this robot { is , are } • the users he mentioned { is , are } • many systems , in addition to VBG a page of free text for each knowledge element , also { permit , permits }
the goal is to predict the correct verb form out of the possible answers in the brackets.
In total, there are approximately 1.4 million test examples in this dataset with varying degrees of difﬁculty.
One proxy to assess the difﬁculty of a test example is through the number of attractors (underlined)— which are deﬁned as intervening nouns of the opposite singular/plural form to the subject.
Each of the example above has zero, one, and four attractors, respectively.
Naturally, examples with fewer numbers of attractors between the head of the syntactic subject and the predicted verb are easier than those with more.
We follow the experimental setup in Linzen et al.
and only use test examples where all the attractors are of contrasting form to the main subject (i.e., all intervening nouns between the subject and the verb must be plural if the subject is singular, and vice versa).
One way to do this task is to train a binary classiﬁer that takes the context and predicts an answer.
We approach this task from a language modeling perspective, where we simply train a language model
[Table]
and take the word with the higher probability between the two possible answers as the prediction.
Success on this task requires a language model that understands syntactic—and in some cases long term—dependencies in natural language.
For a purely sequential memory model to do well on this task, it has to be able to carry dependencies over multiple timesteps and attractors.
On the other hand, our memory augmented recurrent models need to use the random access or stack memory component in conjunction with the sequential memory of their LSTM core to capture these dependencies.
Linzen et al.
(2016) concluded that a vanilla LSTM trained only with a language modeling signal is insufﬁcient for capturing such dependencies.
They reported an overall accuracy of 93.22 with a language modeling objective (using a 50 dimension LSTM), and 99.17 with a supervised binary classiﬁer objective.
We train the best vanilla LSTM (920 dimensions), the best attention-based LSTM ( K = 10 , since the Linzen dataset is derived from Wikipedia articles similar to Wik-2), and the best stack LSTM (multipop stack) on the provided training set that contains sentences of similar structures to the test set.
There are approximately 3 million tokens on the training set ( ∼ 140,000 sentences).
We tune the learning rate and (cid:96) 2 hyperparameter on the development set using perplexity as the tuning criterion.
Our results are shown in Table 2.
We report accuracies per number of attractors, as well as the overall accuracy and perplexity.
Contrary to the Linzen et al.
(2016) results, all of our language models perform surprisingly well on this dataset.
Our vanilla LSTM model outperforms Linzen’s best LSTM by a signiﬁcant margin (99.11 vs. 93.22).
One possible reason is that we are able to train a much bigger LSTM than Linzen—almost 20 times bigger in hidden size.
The results clearly demonstrate the improvements from adding random access and stack memory.
The performance of the attention model slowly degrades to the performance of a vanilla LSTM model as the number of attractors increases, since it becomes more difﬁcult for a random access memory mechanism to attend to the syntactic head in the presence of multiple attractors.
For example, when there are ﬁve attractors, our attention model performs just as well as our vanilla LSTM model.
[Table]
[Table]
the construction of the JJ walls , ﬂoors
[Table]
and VBG walls { is ,are
[Table]
indicates an incorrect prediction, whereas (cid:51) indicates a correct prediction.
In general, we observe that the mistakes made by both the LSTM and attention models that are correctly predicted by the stack model (row 2) typically involve longer sentences regardless of the number of attractors.
LSTMattention stack
The stack model performs best on this dataset, across all numbers of attractors (except zero, tie with attention), Notably, the advantage of the stack model becomes more pronounced as the number of attractors increases.
In §3.4, we analyze how the model uses its stack.
We take this collection of results as evidence that a hierarchical bias introduced by a stack-like data structure helps the language model to learn better syntactic natural language dependencies.
We also investigate whether mistakes are made on the same test examples.
Figure 2 shows a Venn diagram of mistakes made by each of the models.
Most mistakes (6115) are the same across all three models.
It is clear that adding a stack or an attention mechanism improves a vanilla LSTM model, as shown by the signiﬁcant decrease in the number of mistakes that are made only by LSTM (3942) to 1938 and 2534 respectively.
Nonetheless, since there are still a large number of mistakes that are complementary, an interesting future direction is to combine all three kinds of memory models efﬁciently in a single language model.
Table 3 shows examples of mistakes made by each of these models.
3. A NALYSIS
In this section, we analyze how our stack model uses its memory to improve predictions.
Recall that our adaptive continuous stack has 22 possible stack states at each timestep—
STAY 0 , PUSH 0 , STAY 1 , PUSH 1 , . .
. , STAY 10 , PUSH 10 —where the subscript indicates the number of pops that are performed, limited to K = 10 in our experiments.
In Figure 3, we show the number of times each stack state has the highest probability on the Wik-2 test set (left ﬁgure, red denotes STAY and blue denotes PUSH ) and the maximum action probability (right ﬁgure).
The LSTM hidden states are not pushed onto the stack most of the times.
We inspect what is pushed and ﬁnd that the model uses the stack to mostly store hidden states for difﬁcult predictions (more details below).
While the overall distribution of the maximum action probability is not very peaky since the stack is not used
[Figure]
l og f r equen cy                 stay push
[Figure]
max action probability f r equen cy 0.
[Figure] [Figure]
most of
We next investigate how the stack improves accuracy on the Linzen dataset by looking into how it operates when making both a correct prediction and a wrong prediction.
In Figure 4, we show a randomly selected test sentence the loss of basic needs providers VBG from VBN countries {has, have} .
Similar to Figure 3, the red and blue bars represent STAY and PUSH , while the x -axis denotes the number of POP s. The y -axis, on the other hand, denotes the probability of each action.
The green dots above each word shows the magnitude of the norm of the contribution of the memory vector
W h,m m t to the ﬁnal hidden state ˜h t (bigger dots represent higher magnitudes).
Interestingly, the model seems to use and push onto the stack when the next word prediction has a high entropy.
For example, after the word the , of , or from , the model decides to increase its capacity by pushing the current hidden state onto the stack and activates its memory component (as illustrated by bigger green dots).
In §2.1, we discuss connections of our stack model to adaptive computation time and variable computation RNN, which are designed to explicitly increase their capacity for difﬁcult predictions.
Our analysis shows that our stack model also exhibits this kind of behavior.
For the Linzen dataset, we conjecture that the stack model is able to perform the best because it uses the stack as a controller to allow the sequential memory component from its base LSTM to carry longer term dependencies (e.g., tracking the subject of the sentence).
For comparison, we also show how the stack operates when it makes an incorrect prediction in
Figure 5.
Here, the stack behaves similarly, being mostly active for higher entropy predictions, although the model was unable to predict the correct verb right after walls .
construction the of the labyrinth (JJ) walls , ﬂoors , and dividing (VBG) walls {is, are} 
 construction the of the walls , ﬂoors , and walls
[Figure]
C ONCLUSION
We proposed a generalization of the continuous stack model that allows a variable number of pop operations, and compared sequential, random access, and stack memory architectures for recurrent neural network language models.
Our experiments on PTB and Wik-2 showed that adding the stack memory eliminates the need to tune window size in the random access attention model to achieve the best perplexity.
We also evaluated these models on the Linzen syntactic dependencies dataset and demonstrated that the stack augmented model outperforms other methods in terms of both accuracy and perplexity, especially as the number of syntactic attractors increases. A CKNOWLEDGEMENTS
The authors thank Edward Grefenstette for valuable feedback on an earlier draft of this paper, Tal
Linzen for his assistance with the Linzen dataset, and the DeepMind language group for helpful discussions.