RNNs can generate bounded hierarchical languages with optimal memory

Introduction

Recurrent neural networks (RNNs; Elman (1990)) trained on large datasets have demonstrated a grasp of natural language syntax Karpathy et al. (2015); Kuncoro et al. (2018). While considerable empirical work has studied RNN language models’ ability to capture syntactic properties of language Linzen et al. (2016); Marvin and Linzen (2018); Hewitt and Manning (2019); van Schijndel and Linzen (2018), their success is not well-understood theoretically. In this work, we provide theoretical insight into RNNs’ syntactic success, proving that they can efficiently generate a family of bounded hierarchical languages. These languages form the scaffolding of natural language syntax.

Hierarchical structure characterized by long distance, nested dependencies, lies at the foundation of natural language syntax. This motivates, for example, context-free languages Chomsky (1956), a fundamental paradigm for describing natural language syntax. A canonical family of context-free languages (CFLs) is Dyck-$k$, the language of balanced brackets of $k$ types, since any CFL can be constructed via some Dyck-$k$ Chomsky and Schützenberger (1959).

However, while context-free languages like Dyck-$k$ describe arbitrarily deep nesting of hierarchical structure, in practice, natural languages exhibit bounded nesting. This is clear in, e.g., bounded center-embedding Karlsson (2007); Jin et al. (2018) (Figure 1). To reflect this, we introduce and study Dyck-($k$,$m$), which adds a bound $m$ on the number of unclosed open brackets at any time. Informally, the ability to efficiently generate Dyck-($k$,$m$) suggests the foundation of the ability to efficiently generate languages with the syntactic properties of natural language. See §(1.1) for further motivation for bounded hierarchical structure.

In our main contribution, we prove that RNNs are able to generate Dyck-($k$,$m$) as memory-efficiently as any model, up to constant factors. Since Dyck-($k$,$m$) is a regular (finite-state) language, the application of general-purpose RNN constructions trivially proves that RNNs can generate the language Merrill (2019). However, the best construction we are aware of uses $O(k^{\frac{m}{2}})$ hidden units Horne and Hush (1994); Indyk (1995), where $k$ is the vocabulary size and $m$ is the nesting depth, which is exponential.For hard-threshold neural networks, the lower-bound would be $\Omega(k^{\frac{m}{2}})$ if Dyck-($k$,$m$) were an arbitrary regular language. We provide an explicit construction proving that a Simple (Elman; Elman (1990)) RNN can generate any Dyck-($k$,$m$) using only $6m\lceil\log k\rceil-2m=O(m\log k)$ hidden units, an exponential improvement. This is not just a strong result relative to RNNs’ general capacity; we prove that even computationally unbounded models generating Dyck-($k$,$m$) require $\Omega(m\log k)$ hidden units, via a simple communication complexity argument.

Our proofs provide two explicit constructions, one for the Simple RNN and one for the LSTM, which detail how these networks can use their hidden states to simulate stacks in order to efficiently generate Dyck-($k$,$m$). The differences between the constructions exemplify how LSTMs can use exclusively their gates, ignoring their Simple RNN subcomponent entirely, to reduce the memory by a factor of 2 compared to the Simple RNN.We provide implementations at https://github.com/john-hewitt/dyckkm-constructions/.

We prove these results under a theoretical setting that aims to reflect the realistic settings in which RNNs have excelled in NLP. First, we assume finite precision; the value of each hidden unit is represented by $p=O(1)$ bits. This has drastic implications compared to existing work Siegelmann and Sontag (1992); Weiss et al. (2018); Merrill (2019); Merrill et al. (2020). It implies that only regular languages can be generated by any machine with $d$ hidden units, since they can take on only $2^{pd}$ states Korsky and Berwick (2019). This points us to focus on whether languages can be implemented memory-efficiently. Second, we consider networks as language generators, not acceptors;Acceptors consume a whole string and then decide whether the string is in the language; generators must decide which tokens are possible continuations at each timestep. informally, RNNs’ practical successes have been primarily as generators, like in language modeling and machine translation Karpathy et al. (2015); Wu et al. (2016).

Finally, we include a preliminary study in learning Dyck-($k$,$m$) with LSTM LMs from finite samples, finding for a range of $k$ and $m$ that learned LSTM LMs extrapolate well given the hidden sizes predicted by our theory.

In summary, we prove that RNNs are memory optimal in generating a family of bounded hierarchical languages that forms the scaffolding of natural language syntax by describing mechanisms that allow them to do so; this provides theoretical insight into their empirical success.

Hierarchical structure is central to human language production and comprehension, showing up in grammatical constraints and semantic composition, among other properties Chomsky (1956). Agreement between subject and verb in English is an intuitive example:

... The Dyck-$k$ languages—well-nested brackets of $k$ types—are the prototypical languages of hierarchical structure; by the Chomsky-Schützenberger Theorem Chomsky and Schützenberger (1959), they form the scaffolding for any context-free language. They have a simple structure:

..... However, human languages are unlike Dyck-$k$ and other context-free languages in that they exhibit bounded memory requirements. Dyck-$k$ requires storage of an unboundedly long list of open brackets in memory. In human language, as the center-embedding depth grows, comprehension becomes more difficult, like in our example sentence above Miller and Chomsky (1963). Empirically, center-embedding depth of natural language is rarely greater than $3$ Jin et al. (2018); Karlsson (2007). However, it does exhibit long-distance, shallow hierarchical structure:

.. Our Dyck-($k$,$m$) language puts a bound on depth in Dyck-$k$, capturing the long-distance hierarchical structure of natural language as well as its bounded memory requirements.For further motivation, we note that center-embedding directly implies bounded memory requirements in arc-eager left-corner parsers Resnik (1992).

Related Work

This work contributes primarily to the ongoing theoretical characterization of the expressivity of RNNs. Siegelmann and Sontag (1992) proved that RNNs are Turing-complete if provided with infinite precision and unbounded computation time. Recent work in NLP has taken an interest in the expressivity of RNNs under conditions more similar to RNNs’ practical uses, in particular assuming one “unrolling” of the RNN per input token, and precision bounded to be logarithmic in the sequence length. In this setting, Weiss et al. (2018) proved that LSTMs can implement simplified counter automata; the implications of which were explored by Merrill (2019, 2020). In this same regime, Merrill et al. (2020) showed a strict hierarchy of RNN expressivity, proving among other results that RNNs augmented with an external stack Grefenstette et al. (2015) can recognize hierarchical (Context-Free) languages like Dyck-$k$, but LSTMs and RNNs cannot.

Korsky and Berwick (2019) prove that, given infinite precision, RNNs can recognize context-free languages. Their proof construction uses the floating point precision to simulate a stack, e.g., implementing push by dividing the old floating point value by $2$, and pop by multiplying by $2$. This implies that one can recognize any language requiring a bounded stack, like our Dyck-($k$,$m$), by providing the model with precision that scales with stack depth. In contrast, our work assumes that the precision cannot scale with the stack depth (or vocabulary size); in practice, neural networks are used with a fixed precision Hubara et al. (2017).

Our work also connects to empirical studies of what RNNs can learn given finite samples. Considerable evidence has shown that LSTMs can learn languages requiring counters (but Simple RNNs do not) Weiss et al. (2018); Sennhauser and Berwick (2018); Yu et al. (2019); Suzgun et al. (2019), and neither Simple RNNs nor LSTMs can learn Dyck-$k$. In our work, this conclusion is foregone because Dyck-$k$ requires unbounded memory while RNNs have finite memory; we show that LSTMs extrapolate well on Dyck-($k$,$m$), the memory-bounded variant of Dyck-$k$. Once augmented with an external (unbounded) memory, RNNs have been shown to learn hierarchical languages Suzgun et al. (2019); Hao et al. (2018); Grefenstette et al. (2015); Joulin and Mikolov (2015). Finally, considerable study has gone into what RNN LMs learn about natural language syntax Lakretz et al. (2019); Khandelwal et al. (2018); Gulordava et al. (2018); Linzen et al. (2016).

Preliminaries and definitions

A formal language $\mathcal{L}$ is a set of strings $\mathcal{L}\subseteq\Sigma^{*}\omega$ over a fixed vocabulary, $\Sigma$ (with the end denoted by special symbol $\omega$). We denote an arbitrary string as $w_{1:T}\in\Sigma^{*}\omega$, where $T$ is the string length.

The Dyck-$k$ language is the language of nested brackets of $k$ types, and so has $2k$ words in its vocabulary: \Sigma=\{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i},{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}\}_{i\in[k]}. Any string in which brackets are well-nested, i.e., each {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} is closed by its corresponding {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, is in the language. Formally, we can write it as the strings generated by the following context-free grammar:

where $\epsilon$ is the empty string.And $\omega$ appended to the end. The memory necessary to generate any string in Dyck-$k$ is proportional to the number of unclosed open brackets at any time. We can formalize this simply by counting how many more open brackets than close brackets there are at each timestep:

where $\texttt{count}(w_{1:t},a)$ is the number of times $a$ occurs in $w_{1:t}$. We can now define Dyck-($k$,$m$) by combining Dyck-$k$ with a depth bound, as follows: {restatable}[Dyck-($k$,$m$) ]defndefn_dyckkm For any positive integers $k,m$, Dyck-($k$,$m$) is the set of strings

2 Recurrent neural networks

We now provide definitions of recurrent neural networks as probability distributions over strings, and define what it means for an RNN to generate a formal language. We start with the most basic form of RNN we consider, the Simple (Elman) RNN: {restatable}[Simple RNN (generator)]defndefnsimplernn A Simple RNN (generator) with $d$ hidden units is a probability distribution $f_{\theta}$ of the following form:

The Long Short-Term Memory (LSTM) model Hochreiter and Schmidhuber (1997) is a popular extension to the Simple RNN, intended to ease learning by resolving the vanishing gradient problem. In this work, we’re not concerned with learning but with expressivity. We study whether the LSTM’s added complexity enables it to generate Dyck-($k$,$m$) using less memory. {restatable}[LSTM (generator)]defndefnlstmgenerator An LSTM (generator) with $d$ hidden units is a probability distribution $f_{\theta}$ of the following form:

3 Formal language generation

With this definition of RNNs as generators, we now define what it means for an RNN (a distribution) to generate a language (a set). Intuitively, since a formal language is a set of strings $\mathcal{L}$, our definition should be such that a distribution generates $\mathcal{L}$ if its probability mass on the set of all strings $\Sigma^{*}\omega$ is concentrated on the set $\mathcal{L}$. So, we first define the set of strings on which a probability distribution concentrates its mass. The key intuition is to control the local token probabilities $f_{\theta}(w_{t}|w_{1:t-1})$, not the global $f_{\theta}(w_{1:T})$, which must approach zero with sequence length. {restatable}[locally $\epsilon$-truncated support]defndefnepsilontruncated Let $f_{\theta}$ be a probability distribution over $\Sigma^{*}\omega$, with conditional probabilities $f_{\theta}(w_{t}|w_{1:t-1})$. Then the locally $\epsilon$-truncated support of the distribution is the set

This is the set of strings such that the model assigns at least $\epsilon$ probability mass to each token conditioned on the prefix leading up to that token. A distribution generates a language, then, if there exists an $\epsilon$ such that the locally truncated support of the distribution is equal to the language:We also note that any $f_{\theta}$ generates multiple languages, since one can vary the parameter $\epsilon$; for example, any softmax-defined distribution must generate $\Sigma^{*}$ with $\epsilon$ small because they assign positive mass to all strings.

A probability distribution $f_{\theta}$ over $\Sigma^{*}$ generates a language $\mathcal{L}\subseteq\Sigma^{*}$ if there exists $\epsilon>0$ such that the locally $\epsilon$-truncated support of $f_{\theta}$ is $\mathcal{L}$.

Formal results

For our results, we first present two theorems for the Simple RNN and LSTM that use $O(mk)$ hidden units by simulating a stack of $m$ $O(k)$-dimensional vectors, which are useful for discussing the constructions. Then we show how to reduce to $O(m\log k)$ via an efficient encoding of $k$ symbols in $O(\log k)$ space.

Using the same mechanisms as in the proofs above but using an efficient encoding of each stack element in $O(\log k)$ units, we achieve the following.

While we have emphasized expressive power under memory constraints—what functions can be expressed, not what is learned in practice—neural networks are frequently intentionally overparameterized to aid learning Zhang et al. (2017); Shwartz-Ziv and Tishby (2017). Even so, known constructions for Dyck-($k$,$m$) would require a number of hidden units far beyond practicality. Consider if we were to use a vocabulary size of $100{,}000$, and a practical depth bound of $3$. Then if we were using a $k^{m+1}$ hidden unit construction to generate Dyck-($k$,$m$), we would need $100{,}000^{4}=10^{20}$ hidden units. By using our LSTM construction, however, we would need only $3\times 3\times\lceil\log_{2}(100{,}000)\rceil-1\times 3=150$ hidden units, suggesting that networks of the size commonly used in practice are large enough to learn these languages.

Lower bound.

Stack constructions in Simple RNNs

The memory needed to close all the brackets in a Dyck-($k$,$m$) prefix $w_{1:t}$ can be represented as a stack of (yet unclosed) open brackets [{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}], $m^{\prime}\leq m$; reading each new parenthesis either pushes or pops from this stack. Informally, all of our efficient RNN constructions generate Dyck-($k$,$m$) by writing to and reading from an implicit stack that they encode in their hidden states. In this section, we present some challenges in a naive approach, and then describe a construction to solve these challenges. We provide only the high-level intuition; rigorous proofs are provided in the Appendix.

We start by describing what is achievable with an extended model family, second-order RNNs Rabusseau et al. (2019); Lee et al. (1986), which allow their recurrent matrix $W$ to be chosen as a function of the input (unlike any of the RNNs we consider.) Under such a model, we show how to store a stack of up to $m$ $k$-dimensional vectors in $mk$ memory. Such a representation can be thought of as the concatenation of $m$ $k$-dimensional vectors in the hidden state, like this:

We call each $k$-dimensional component a “stack slot”. If we want the first stack slot to always represent the top of the stack, then there’s a natural way to implement pop and push operations. In a push, we want to shift all the slots toward the bottom, so there’s room at the top for a new element. We can do this with an off-diagonal matrix $W_{\text{push}}$:Note that only needing to store $m$ things means that when we push, there should be nothing in slot $m$; otherwise, we’d be pushing element $m+1$.

This would implement the $Wh_{t-1}$ part of the Simple RNN equation. We can then write the new element (given by $Ux_{t}$) to the first slot. If we wanted to pop, we could do so with another off-diagonal matrix $W_{\text{pop}}$, shifting everything towards the top to get rid of the top element:

This won’t work for a Simple RNN because it only has one $W$.

2 A Simple RNN Stack in 2​m​k2𝑚𝑘2mk memory

Our construction gets around the limitation of only having a single $W$ matrix in the Simple RNN by doubling the space to $2mk$. Splitting the space $h$ into two $mk$-sized partitions, we call one $h_{\text{pop}}$, the place where we write the stack if we see a pop operation, and the other $h_{\text{push}}$ analogously for the push operation. If one of $h_{\text{pop}}$ or $h_{\text{push}}$ is empty (equal to ) at any time, we can try reading from both of them, as follows:

Our $W$ matrix is actually the concatenation of two of the $W_{\text{pop}}$ and $W_{\text{push}}$ matrices. Now we have two candidates, both $h_{\text{push}}$ and $h_{\text{pop}}$; but we only want the one that corresponds to push if w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, or pop if w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}. We can zero out the unobserved option using the term $Ux_{t}$, adding a large negative value to every hidden unit in the stack that doesn’t correspond to push if $x_{t}$ is an open bracket {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, or pop if $x_{t}$ is a close bracket {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}:

So Ux_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}}=[e_{i};0,\dots,-\beta-1,\dots], where $e_{i}$ is a one-hot representation of {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, and Ux_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}}=[-\beta-1,\dots,-\beta-1,0,\dots], where $\beta$ is the very large value we specified in our finite-precision arithmetic. Thus, when we apply the sigmoid to the new $Wh_{t-1}+Ux_{t}+b$, whichever of $\{h_{\text{pop}},h_{\text{push}}\}$ doesn’t correspond to the true new state is zeroed out.For whichever of $h_{\text{tmp}}\in\{h_{\text{pop}},h_{\text{push}}\}$ that is not zeroed out, $\sigma(h_{\text{tmp}})\not=h_{\text{tmp}}$. Hence, we scale all of $U$ and $W$ to be large, such that $\sigma(Wh_{t-1})\in\{0,1\}$.

Stack construction in LSTMs

We could implement our $2mk$ Simple RNN construction in an LSTM, but its gating functions suggest more flexibility in memory management, and Levy et al. (2018) claim that the LSTM’s modeling power stems from its gates. With an LSTM, we achieve the $mk$ of the oracle we described, all while exclusively using its gates.

To implement a stack using the LSTM’s gates, we use the same intuitive description of the stack as before: $m$ stack slots, each of dimensionality $k$. However while the top of the stack is the first slot in the Simple RNN, the bottom of the stack is the first slot in the LSTM. Before we discuss mechanics, we introduce the memory dynamics of the model. Working through the example in Figure 2, when we push the first open bracket, it’s assigned to slot $1$; then the second and third open brackets are assigned to slots $2$ and $3$. Then a close bracket is seen, so slot $3$ is erased. In general, the stack is represented in a contiguous sequence of slots, where the first slot represents the bottom of the stack. Thus, the top of the stack could be at any of the $m$ stack slots. So to allow for ease of linearly reading out information from the stack, we store the full stack only in the cell state $c_{t}$, and let only the slot corresponding to the top of the stack, which we’ll refer to as the top slot, through the output gate to the hidden state $h_{t}$.

pop mechanics.

To implement a pop operation, the forget gate finds the top slot, which is the slot farthest from slot $1$ that isn’t empty (that is, that encodes some {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}.) In practice, we do this by guaranteeing that the forget gate is equal to $1$ for all stack slots before (but excluding) the last non-empty stack slot. Since this last non-empty stack slot encodes the top of the stac,, for it and all subsequent (empty) stack slots, the forget gate is set to .The input gate and the output gate are both set to $\mathbf{0}$. This erases the element at the top of the stack, summarized in the following diagram:

output mechanics.

We’ve so far described how the new cell state $c_{t}$ is determined. So that it’s easy to tell what symbol is at the top of the stack, we only want the top slot of the stack passing through the output gate. We do this by guaranteeing that the output gate is equal to for all stack slots from the first slot to the top slot (exclusive). The output gate is then set to $1$ for this top slot (and all subsequent empty slots,) summarized in the following diagram:

Defining the generating distribution

So far, we’ve discussed how to implement implicit stack-like memory in the Simple RNN and the LSTM. However, the formal claims we make center around RNNs’ ability to generate Dyck-($k$,$m$).

Assume that at any timestep $t$, our stack mechanism has correctly pushed and popped each open and close bracket, as encoded in $h_{t},c_{t}$. We still need to prove that our probability distribution,

assigns greater than $\epsilon$ probability only to symbols that constitute continuations of some string in Dyck-($k$,$m$), by specifying the parameters $V$ and $b_{v}$.

If and only if fewer than $m$ elements are on the stack, all {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} must be assigned $\geq\epsilon$ probability. This encodes the depth bound. In our constructions, $m$ elements are on the stack if and only if stack slot $m$ is non-zero. So, each row V_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}} is zeros except for slot $m$, where each dimension is a large negative number, while the bias term $b_{v}$ is positive.

Observing the end of the string ω𝜔\omega.

If and only if elements are on the stack, the string can end. The row $V_{\omega}$ detects if any stack slot is non-empty.In particular, the bias term $b_{\omega}$ is positive, but the sum $V_{\omega}h_{t-1}+b_{\omega}$ is negative if the stack is not empty.

The close bracket {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i} can be observed if and only if the top stack slot encodes {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}. Both the Simple RNN construction and the LSTM construction make it clear which stack slot encodes the top of the stack. In the Simple RNN, it’s always the first slot. In the LSTM, it’s the only non-empty slot in $h_{t}$. In our stack constructions, we assumed each stack slot is a $k$-dimensional one-hot vector $e_{i}$ to encode symbol {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}. So in the Simple RNN, V_{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}{}_{i} reads the top of the stack through a one-hot vector $e_{i}$ in slot $1$, while in the LSTM it does so through $e_{i}$ in all $m$ slots. This ensures that V_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}}^{\top}h_{t} is positive if and only if {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} is at the top of the stack.

2 Generation in O​(m​log⁡k)𝑂𝑚𝑘O(m\log k) memory

Thus, we can always detect which symbol {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} is encoded by setting b_{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}{}_{i}=\log k-0.5.In actuality, we use a slightly less compact encoding, spending $\log k-1$ more hidden units set to $1$, to incrementally subtract $\log k-1$ from all logits. Then the bias terms b_{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}{}_{i} are set to $0.5$, avoiding possible precision issues with representing the float $\lceil\log k\rceil-0.5$.

Experiments

Our proofs have concerned constructing RNNs that generate Dyck-($k$,$m$); now we present a short study connecting our theory to learning Dyck-($k$,$m$) from finite samples. In particular, for $k\in\{2,8,32,128\}$ and $m\in\{3,5\}$, we use our theory to set the hidden dimensionality of LSTMs to $3m\lceil\log k\rceil-m$, and train them as LMs on samples from a distributionDefined in Appendix H over Dyck-($k$,$m$). For space, we provide an overview of the experiments, with details in the Appendix. We evaluate the models’ abilities to extrapolate to unseen lengths by setting a maximum length of $84$ for $m=3$, and $180$ for $m=5$, and testing on sequences longer than those seen at training time.Up to twice as long as the training maximum. For our evaluation metric, let $p_{j}$ be the probability that the model predicts the correct closing bracket given that $j$ tokens separate it from its open bracket. We report $\text{mean}_{j}p_{j}$, to evaluate the model’s bracket-closing memory.

For all configurations, we find that the LSTMs using our memory limit achieve error less than $10^{-4}$ when trained on 20 million tokens. Strikingly, this is despite the fact that for large $m$ and $k$, a small fraction of the possible stack states is seen at training time;See Table 3. this shows that the LSTMs are not simply learning $k^{m}$ structureless DFA states. Learning curves are provided in Figure 3.

Discussion and conclusion

We proved that finite-precision RNNs can generate Dyck-($k$,$m$), a canonical family of bounded-depth hierarchical languages, in $O(m\log k)$ memory, a result we also prove is tight. Our constructions provide insight into the mechanisms that RNNs and LSTMs can implement.

The Chomsky hierarchy puts all finite memory languages in the single category of regular languages. But humans generating natural language have finite memory, and context-free languages are known to be both too expressive and not expressive enough Chomsky (1959); Joshi et al. (1990). We thus suggest the further study of what structure networks can encode in their memory (here, stack-like) as opposed to (just) their position in the Chomsky hierarchy. While we have settled the representation question for Dyck-($k$,$m$), many open questions still remain: What broader class of bounded hierarchical languages can RNNs efficiently generate? Our experiments point towards learnability; what class of memory-bounded languages are efficiently learnable? We hope that answers to these questions will not just demystify the empirical success of RNNs but ultimately drive new methodological improvements as well.

Code for running our experiments is available at https://github.com/john-hewitt/dyckkm-learning. An executable version of the experiments in this paper is on CodaLab at https://worksheets.codalab.org/worksheets/0xd668cf62e9e0499089626e45affee864.

Acknowledgements

The authors would like to thank Nelson Liu, Amita Kamath, Robin Jia, Sidd Karamcheti, and Ben Newman. JH was supported by an NSF Graduate Research Fellowship, under grant number DGE-1656518. Other funding was provided by a PECASE award.

References

Appendix A Appendix outline

This Appendix has the following order. In (§B), we provide a definition of Dyck-($k$,$m$) equivalent to that in the main text but more useful for our proofs. In (§C), we state preliminary definitions and assumptions, and prove the lower-bound of $\Omega(m\log k)$ hidden units to generate Dyck-($k$,$m$). In (§D), we formally introduce our Simple RNN stack construction, and prove its correctness in a lemma. In (§E), we formally introduce our LSTM stack construction, and prove its correctness in a lemma. In (§F), we prove that a linear (+softmax) decoder on the Simple RNN and LSTM hidden states can be used to generate Dyck-($k$,$m$) in $O(mk)$ hidden units using a $1$-hot encoding of stack elements. In this section we also prove that a general RNN construction of DFAs allows for generation of Dyck-($k$,$m$) in $O(k^{m+1})$ hidden units. In (§G), we provide an alternative encoding of elements in our stack constructions for the Simple RNN and LSTM that uses $O(\log k)$ space per element, and prove that our stack constructions still hold using this encoding. We provide a linear (+softmax) decoder on the hidden states of the Simple RNN and LSTM (when using the $O(\log k)$ stack element encoding) that can be used as a drop-in replacement for the decoder from the $1$-hot representations, thus generating Dyck-($k$,$m$). This proving that both the Simple RNN and LSTM generate Dyck-($k$,$m$) in $O(m\log k)$ hidden units.

Appendix B The Dyck-(k𝑘k,m𝑚m) languages

To better understand the success of neural networks on natural language syntax, we aim for a formal language that models the unbounded recursiveness of natural language while also reflecting its bounded memory requirements. We thus introduce the Dyck-($k$,$m$) languages, corresponding to sequences of balanced brackets of $k$ types with a maximal number of $m$ unclosed brackets at any point in the sequences (yielding a bound stack depth of $m$ to parse such sentences).

Though in the main text we defined Dyck-($k$,$m$) by intersecting Dyck-$k$ with a language that simply bounds the difference between the number of open brackets and the number of close brackets, here we provide an equivalent definition that will aid in our proofs. Each Dyck-($k$,$m$) language, specified by fixing a value of $m$ and $k$, is defined using a deterministic finite automaton (DFA). Here, we provide a general description of any Dyck-($k$,$m$) DFA.

Formally, we define each language by the deterministic finite automaton $\mathcal{D}_{m,k}=(Q,\Sigma,\delta,F,q_{0})$. The vocabulary $\Sigma$ consists of $k$ types of open brackets: \{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}\}_{i=1,\dots,k} corresponding closing brackets \{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}\}_{i=1,\dots,k}. Strings over the vocabulary are $w_{1:T}\in\Sigma^{*}\omega$, where $\omega\not\in\Sigma$ is a special symbol representing the end of the sequence. $\Sigma\cup\{\omega\}$ are collectively referred to as symbols.

This slightly nonstandard requirement allows for a natural connection with language models, which must estimate the probability that a string ends at any given token.But crucially, since $\omega\not\in\Sigma$, the language model need not define a distribution after $\omega$ is seen; it can only be the last token. Overloading notation, we’ll also use $\mathcal{D}_{m,k}$ to refer to the language (the set of strings) itself, defined as the strings accepted by the DFA.

We now define the states $q\in Q$. First, we define reject state $r$, and accept state $[\omega]$. Each other state is uniquely identified by a list of open bracket symbols of length up to $m$; thus the full set of states is provided by:

The number of states is thus $k^{m+1}+1$, where all but two states reflect a list of open brackets. We will denote each list [{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}}{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{2}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}], $0\leq m^{\prime}\leq m$ as a stack state with $m^{\prime}$ elements, and the value of {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}} as the top element of the stack. We let q_{0}=[]=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}}{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{2}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}], $m^{\prime}=0$, the empty stack state.

B.2 Transition function, δ𝛿\delta

We now define the transition function, $\delta$.

The state $[]$ can transition either to the accept state or to another stack state:

while any other symbol transitions to the reject state.

For any state of the form [{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}}{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{2}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}], where $m^{\prime}<m$, an open bracket can be pushed to the list (since $m^{\prime}<m$), or the last open bracket can be removed, by observing its corresponding close bracket.

All other symbols transition to the reject state.

For any state of the form [{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}}{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{2}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m}}], that is, the list is of length $m$ and thus full, the close bracket of the top of the stack removes that bracket,

All other symbols transition to the reject state.

The reject state $r$ transitions to itself for all symbols.

No transitions from the accept state need be defined, since only $\delta([],\omega)=[\omega]$, and $\omega$ must be the last symbol of any string in the universe, and can only occur once.

This accounts for the transition function from all states in $Q$, and completes our definition of the DFA $D_{m,k}$. We are now prepared to formally define the language:

Overloading notation, for any $w_{1:T}\in\Sigma^{*}\omega$, for any $t\in 1,\dots,T$, we’ll say $q_{t}=D_{m,k}(w_{1:t})$, denoting the state that $D_{m,k}$ is in after consuming $w_{1:t}$.

Appendix C Preliminaries

* Through the notion of $\epsilon$-truncated support, a language model specifies which tokens are allowable continuations of each string prefix by assigning them greater than $\epsilon$ probability. With this, we’re ready to connect RNNs and formal languages:

A probability distribution $f_{\theta}$ generates a formal language $\mathcal{L}$ if there exists $\epsilon>0$ such that the $\epsilon$-truncated support $\mathcal{L}_{f_{\theta}}$ is equal to $\mathcal{L}$.

Under our finite-precision setting, a reasonable assumption about the properties of the sigmoid and hyperbolic tangent functions make our claim considerably simpler. The sigmoid function, $\sigma(x)=\frac{1}{1+e^{-x}}$, has range $(0,1)$, excluding its boundaries $\{0,1\}$. However, in a finite-precision arithmetic, $\sigma(x)$ cannot become arbitrarily close to or $1$. In fact, in popular deep learning library PyTorchUnder the float datatype, PyTorch v1.3.1 Paszke et al. (2019)., $\sigma(x)$ is exactly equal to $1$ for all $x>6$. We define the floating point sigmoid function to equal or $1$ if the absolute value of its input is greater or equal in absolute value to some threshold $\beta$:

Similarly for the hyperbolic tangent function, we define:

For the rest of this paper, we will refer to $\sigma_{\text{fp}}$ as $\sigma$, and $\text{tanh}_{\text{fp}}$ as tanh.

C.2 Lower bound of Ω​(m​log⁡k)Ω𝑚𝑘\Omega(m\log k) for generation of Dyck-(k𝑘k,m𝑚m)

We provide a communication complexity argument. Assume for contradiction that there exists such a machine with $d<\frac{m\log k}{p}$. Consider any string w_{1:m}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m}} of $m$ open brackets, where each is one of $k$ types. There are $k^{m}$ such strings, but because $d<\frac{m\log k}{p}$, we have that the model has $2^{dp}<k^{m}$ possible representations, so at least two such strings must share the same representation. Let such a pair be $w\not=w^{\prime}$, where likewise $w^{\prime}$ is defined as {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i^{\prime}_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i^{\prime}_{m}}. Consider the sequence of open bracket indices that defines $w$: $i_{1},\dots,i_{m}$. Let $s$ be the string {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i_{m}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i_{1}} of close brackets in the reverse order of the indices of $w$, and likewise $s^{\prime}$ for $w^{\prime}$.

The string $w::s$, where $::$ denotes concatenation, is in Dyck-($k$,$m$). However, $w^{\prime}::s$ is not in Dyck-($k$,$m$) as it breaks the well-nested brackets condition wherever {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{j}}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i^{\prime}_{j}}. This must occur at least once, since $w\not=w^{\prime}$. Since the model assigns them the same representation even though they must be distinguished to generate Dyck-($k$,$m$), the model does not generate Dyck-($k$,$m$).

Appendix D Proving Simple RNN stack correspondence in 2​m​k2𝑚𝑘2mk units

In this section, we provide a formal description of our Simple RNN stack construction, and introduce and prove the stack correspondence lemma, to guarantee its correctness. \defnsimplernn*

where $[\dots]$ denotes concatenation of vectors into a $km$-dimensional vector. We note that $R$ is injective, that is, the inverse map $R^{-1}$ is well-defined on all vectors in the image of $R$.

Later, when constructing the efficient $O(m\operatorname{log}k)$ memory encoding in Section G.1, we will swap the one-hot vectors $e_{i}$ for other vectors that will also have entries in $\{0,1\}$.Looking ahead to the more complex LSTM construction, we’ll introduce notation to denote these vectors e_{i}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}), and look to replace $\psi$ when we develop the $O(m\log k)$ construction.

If $h=\mathcal{S}(\sigma,q)$, then the definition guarantees that $\mathcal{Q}(h)=q$.

Defining Transition Matrices

To define the RNN (apart from $V$, $b_{v}$), we have to specify the matrices $W,U,E,$, and the vector $b$.

To specify the matrices $E$ and $U$, we need to fix an assignment to the integers $1,\dots,2k$ to the symbols in $\Sigma$. We will assign the integers $1,\dots,k$ to the opening brackets (i.e., {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{1}\mapsto 1,\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{k}\mapsto k), and the integers $1,\dots,k$ to the closing brackets (i.e., {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{1}\mapsto k+1,\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{k}\mapsto 2k).

Here, we write $I_{k\times k}$ for the identity matrix, and ${\bf 1}_{k\times k}$ for a matrix filled entirely with ones.

Proof of Correctness

To prove correctness of the construction, the first step is to show that the transition dynamics of the Simple RNN correctly simulates the dynamics of the stack when consuming one symbol.

Let \phi({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})=push, \phi({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i})=pop.

Assume that the hidden state $h_{t}$ encodes a stack $q$:

Let $w_{t+1}$ be the new input symbol, and assume that $w_{1:t+1}$ is a valid prefix of a word in Dyck-($k$,$m$). Let $h_{t+1}$ be the next hidden state, after reading $w_{t+1}$. Then

Before showing this lemma, we note the following property of the activation function $\sigma$, under the finite precision assumption:

We can write $h_{t}=[h_{push},h_{pop}]$. Set $h^{\prime}=h_{push}+h_{pop}$. By construction of $\mathcal{S}$,

At this point, we note that the only relevant property of the encodings $h^{\prime}_{i}$ for this proof is that their entries are in $\{0,1\}$, not that they are one-hot vectors. This will make it possible to plug in a more efficient $O(\log k)$ encoding later in Section G.1.

Case 1:

Assume w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}. We have

which is equal to S(push,\delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})).

Case 2:

Assume w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}. We have

which is equal to S(push,\delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i})). ∎

From the previous lemma, we can derive the following Stack Correspondence lemma, which asserts that the Simple RNN correctly reflects stack dynamics over an entire input string:

For all strings $w_{1:T}$ in Dyck-($k$,$m$), for all $t=1,...,T$, let $q_{t}$ be the state of $D_{m,k}$ after consuming prefix $w_{1:t}$. Let $h_{t}$ be the RNN hidden state after consuming $w_{1:t}$. Then if $q_{t}\not\in\{[\omega]\}$,

The claim is shown by induction over $t$, applying the previous lemma in each step. We will show the claim for all $t=0,1,\dots,T$, setting $h_{0}$ to be the zero vector ${\bf 0}_{2km}$.

As $\mathcal{Q}(h_{0})=q_{0}$, this proves the claim for $t=0$. To prove the inductive step, we assume $\mathcal{Q}(h_{t})=q_{t}$ has already been shown.

Noting $q_{t+1}=\delta(q_{t},w_{t+1})$, we obtain

As noted in the definition of $\mathcal{Q}$ above (D.3), this entails

Appendix E Proving LSTM stack correspondence in m​k𝑚𝑘mk units

In this section, we prove our lemma describing the stack construction implemented by an LSTM. To do so, we first define the LSTM: \defnlstmgenerator*

First, we provide notation for describing the structure of the memory cell ($c_{t}$) that $f_{\theta}$ will maintain. Next, we proceed to describe precisely, but without referring to the LSTM equations, how this memory changes over time. We then describe how these high-level dynamics are implemented by the LSTM equations, assuming that we’re able to set the values of the intermediate values (gates and new cell candidate) as desired. We then explicitly construct LSTM weight matrices that provide the desired intermediate values. This completes the proof of the memory dynamics, which we formalize in a stack correspondence lemma once we’ve introduced the notation.

Applying this function to each stack slot allows us to map from $s_{t,1},\dots,s_{t,m}$ to stack states:

where $m^{\prime}\leq m$ is the maximal integer such that $\psi(s_{t,m^{\prime}})\not=\varnothing$. Intuitively, this means we filter out any empty stack slots at the end of the sequence (those slots $s_{t,j}$ where $j>m^{\prime}$.) We’ll also make use of the inverse, $\mathcal{Q}^{-1}$, mapping DFA stack states to stack slot lists,

where slots $s_{m^{\prime}+1},\dots,s_{m}$ are the zero vector $\mathbf{0}$ because only the first $m^{\prime}$ slots encode a symbol {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, and $\psi(\mathbf{0})=\varnothing$.

These one-hot encodings ($e_{i}$) of symbols ({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}) allow for clear exposition, but are not fundamental to the construction; in (§ G), we replace these $k$-dimensional one-hot encodings given by $\psi$ with $O(\log k)$-dimensional encodings. Throughout our description of the stack construction, we’ll explicitly rely on the following property, so we can easily replace this particular $\psi$ in (§ G) by providing another $\psi^{\prime}$ with the same property.

that is, the sum of all dimensions of the representation is equal to $1$, and

that is, the representation takes on values only in and $1$, and

that is, the empty stack slot is encoded by the zero vector, and

that is, encodings of symbols {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} are unique, and none are equal to to the encoding of the empty symbol $\varnothing$.

The encoding we’ve so far provided, \psi^{-1}(e_{i})={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, obeys these properties: the sum of a single one-hot vector is $1$, no two such vectors are the same, one-hot vectors only take on values in $\{0,1\}$, and we let $\psi(\mathbf{0})=\varnothing$.

E.2 Description of stack state dynamics

Before we describe the memory dynamics, we introduce a useful property of the stack slots.

A cell state $c_{t}$ composed of stack slots $s_{t,1},\dots,s_{t,m}$ is $j$-top if there exists $i\in\{1,\dots,k\}$ such that s_{t,j}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}), and for all $j^{\prime}\in\{j+1,\dots,m\}$, $s_{t,j^{\prime}}=\mathbf{0}$.

Intuitively, we’ll enforce the constraint that a cell state that is $j$-top encodes the element {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{m^{\prime}} at the top of the stack [{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{m^{\prime}}] in stack slot $j=m^{\prime}$.

We’re now ready to describe the memory dynamics of $f_{\theta}$. The memory $c_{t}$ is initialized to $\mathbf{0}$. So for all $j\in\{1,\dots,m\}$,

The recurrent equation of stack slots is then

Intuitively, the first case of Equation E.10 implements pop, the second case implements push (if pushing to an otherwise empty stack) and the third case implement push (if pushing to a non-empty stack). The fourth case specifies that slots that are not pushed to or popped are maintained from timestep $t-1$ to timestep $t$.

Intuitively, this means only the top of the stack is stored in the hidden state, and it’s stored in whichever units correspond to the $j$-top slot in $c_{t}$. Since \psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}) takes on values in $\{0,1\}$, \text{tanh}(\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})) takes on values in $\{0,\text{tanh}(1)\}$.

We’re now ready to formalize the first part of our proof in a lemma that encompasses how $f_{\theta}$ structures and manipulates its memory.

For all strings $w_{1:T}$ in Dyck-($k$,$m$), for all $t=1,...,T$, let $q_{t}$ be the state of $D_{m,k}$ after consuming prefix $w_{1:t}$. Let $c_{t}$ be the hidden state of the LSTM after consuming $w_{1:t}$. Then if $q_{t}=\not\in\{[\omega]\}$,

and letting q_{t}=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}] for $m^{\prime}\leq m$ without loss of generality,

where if $q_{t}=[]$, then $m^{\prime}=0$ and $h_{t,j}=\mathbf{0}$ for all $j$.

Note again that the value of $\psi^{-1}$ is not specified by the lemma; whereas so far we’ve used \psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})=e_{i}, a one-hot encoding, in (§ G) we’ll replace $\psi^{-1}$ while maintaining the lemma.

Intuitively, Equation E.12 states that $f_{\theta}$ keeps an exact representation of the stack of unclosed open brackets so far, a statement solely concerning the cell state $c_{t}$. However, all intermediate values in the LSTM equations are functions of the hidden state $h_{t}$, not the cell state directly; Equation E.13 states how the top of the stack is represented in the hidden state, $h_{t}$. We state these as one lemma since it is convenient to prove them by induction together.

E.3 Proof of the stack correspondence lemma assuming stack slot dynamics

We first prove Lemma 4 under the assumption that the stack slot dynamics given in Equation E.10 hold. We proceed by induction on the prefix length.

When the prefix length is , the DFA is in state $[]$, and $s_{0,j}=\mathbf{0}$ for all $j$, so $\mathcal{Q}(s_{0,1},\dots,s_{0,m})=[]$, as required. Further, we have $h_{0}=\mathbf{0}$, as required. This completes the base case.

Assume for the sake of induction that Lemma 4 holds for strings up to length $t$. Now consider a prefix $w_{1:t+1}$. From it, we can take the prefix $w_{1:t}$. Running $f_{\theta}$, we have $s_{t,1},\dots,s_{t,m}$. By the induction hypothesis, we have that \mathcal{Q}(s_{t,1},\dots,s_{t,m})=q_{t}=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}], (Equation E.12) as well as that

by Equation E.13. Now, the symbol $w_{t+1}$ can be one of $2k+1$ symbols: any of the $k$ open brackets {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, the $k$ close brackets {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, or $\omega$. The proof now proceeds by each of these three cases.

Without loss of generality, we let q_{t}=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}] for $m^{\prime}<m$. Strict inequality is guaranteed since if $m^{\prime}=m$, then \delta(q_{t},{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})=r, meaning $w_{1:t+1}$ cannot be a prefix of any string in Dyck-($k$,$m$), contradicting a premise of the lemma. Through the inverse of the stack mapping Equation E.3, we have

the stack slots of $f_{\theta}$ at timestep $t$. Since s_{t,m^{\prime}}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{j}}) and $s_{t,m^{\prime}+1}=\mathbf{0}$, we have that $c_{t}$ is $m^{\prime}$-top. Thus, since w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, the second or third condition in Equation E.10 (depending on whether $m^{\prime}=0$) dictate that s_{t+1,m^{\prime}+1}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}) (where the $i$ is because w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}). All other stack slots fall under the condition $s_{t+1,j}=s_{t,j}$.

With this, we can reason about the DFA state encoded by $c_{t+1}$ at timestep $w_{1:t+1}$:

which is the DFA’s state at timestep $t+1$, as required.

Finally, we reason about $h_{t+1,j}$. We’ve just shown that that $c_{t+1}$ is $(m^{\prime}+1)$-top, and that s_{t+1,m^{\prime}+1}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}). By Equation E.11, we have that h_{t+1,m^{\prime}+1}=\text{tanh}(\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})), and $h_{t+1,j}=\mathbf{0}$ for all $j\not=m^{\prime}+1$ as required, completing this case.

We have q_{t}=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},\dots,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}] for $0<m^{\prime}$. this inequality is strict since if $m^{\prime}=0$, then $q_{t}=[]$, and \delta(q_{t},{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i})=r, meaning $w_{1:t+1}$ is not a prefix of any string in $D_{m,k}$. Thus, we have that $c_{t}$ is $m^{\prime}$-top. Since w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, we have that $s_{t+1,m^{\prime}}=\mathbf{0}$, and $s_{t+1,j}=s_{t,j}$ for all $j\not=m^{\prime}$. Thus, we have

and thus the DFA state corresponding to the stack slots is:

Which is \delta(q_{t},{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}), as required. By the same reasoning as for case w_{t+1}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, since $c_{t+1}$ is $(m^{\prime}-1)$-top and that s_{t+1,m^{\prime}-1}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{m^{\prime}-1}), we have that h_{t+1,m^{\prime}-1}=\text{tanh}(\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})) and $h_{t+1,j}=\mathbf{0}$ for all $j\not=m^{\prime}-1$, as required. This completes the case.

𝑡1𝜔w_{t+1}=\omega: If $w_{t+1}=\omega$, then by the definition of $\delta$, $\delta(q_{t},\omega)\in\{[\omega],r\}$; hence the premises of Lemma 4 don’t hold, so the lemma vacuously holds.

Summary

In this section we’ve proved the stack correspondence lemma for our LSTM construction, assuming we can implement the dynamics of the model as specified in Equations E.10 E.11; we have yet to show that the LSTM update can implement the equations as promised.

E.4 Implementation of stack state dynamics in LSTM equations assuming intermediate values

In this subsection, we define how the dynamics defined in Equations E.10, E.11 are implemented in the LSTM equations, referring to intermediate gate values, assuming such values can be reached with some setting of LSTM parameters.

How these fulfill Equations E.10, E.11? We can set the gate values to do so as follows. First, the forget gate is used to detect the condition under which slot $j$ should be erased:

Second, the new cell candidate expression is used to give the option of writing a new open bracket, $\psi^{-1}(w_{t})$ to any stack slot:

Third, the input gate determines which stack slot is the first non-empty slot in the sequence, and only allows the information of the new open bracket to be written to that slot:

Finally, the output gate is set to identify the top of the stack and not let any other non-$\mathbf{0}$ (i.e., all stack slots below the top) through:

Many of these conditions refer to values, like $c_{t}$ and $c_{t-1}$, not available when computing the gates (as only $h_{t-1}$ and $x_{t}$ are available); it is convenient to refer to these conditions and later show how they are implemented with the available values.

E.5 Proof of stack state dynamics given gate and new cell candidate values

In this section, we prove that Equations E.10, E.11 are implemented by the LSTM $f_{\theta}$ given the gate and new cell candidate values defined in the previous subsection.

We start with Equation E.10, the definition of the stack slot dynamics, proceeding by each of the four cases for defining $s_{t,j}$ in Equation E.10.

where the last equality hold because of the condition that $c_{t-1}$ is $(j-1)$-top, which implies $s_{t-1,j}=\mathbf{0}$.

Proving Equation E.11 is implemented

Next, we prove that Equation E.11 is implemented by the LSTM $f_{\theta}$, given that Equation E.10 is. First, we must show $h_{t,j}$ is equal to $\text{tanh}(e_{i})$ if $c_{t}$ is $j$-top. If $c_{t}$ is $j$-top, it must be that $s_{t,j}=e_{i}$ for some $i$. Since $c_{t}$ is $j$-top, we have $o_{t,j}=\mathbf{1}$, and $s_{t,j}=e_{i}$. Plugging these values into the hidden state expression in Equation E.15, we get:

If that condition does not hold, we must have that $c_{t}$ not $j$-top. thus, $c_{t}$ is either $j$-top for some $j^{\prime}<j$ or $j^{\prime}>j$, or $c_{t}=\mathbf{0}$.

If $c_{t}$ is $j^{\prime}$-top for $j^{\prime}<j$, then we have that $s_{t,j}=\mathbf{0}$, by the definition of $j$-top. In this case, we have

If $c_{t}$ is $j^{\prime}$-top for $j^{\prime}>j$, then we have that $o_{t,j}=\mathbf{0}$, by Equation E.19. In this case, we have

Finally, if $c_{t}=\mathbf{0}$, meaning it is not $j$-top for any $j$, then

Summary

So far, we’ve proved the stack correspondence lemma’s induction step for LSTMs assuming that we can provide parameters of the network such that the values assumed in Equations E.16, E.17, E.18, E.19, that is, the gates and new cell candidate, are achieved. We have yet to provide the settings of parameters to do so.

E.6 Construction and proof of LSTM parameters providing gate and new cell candidate values

The equation defining $o_{t,j}$ is as follows:

and likewise for the second condition. The third condition relies on the fact that $\psi^{-1}(\varnothing)=\mathbf{0}$.

that is, each row of $U_{f}$ is equal to the sum, over all close brackets, of the embedding of that close bracket (transposed.) Since all x_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}} are orthogonal and unit norm (by the definition of $E$), this gives the following intermediate value:

that is, all units of the bias are equal to the same value.

As we stated earlier, Equation E.19 refers to conditions on $c_{t}$, which do not show up in the computation of $o_{t}$; instead, $o_{t}$ is a function of $h_{t-1}$, and $x_{t}$. As such, we re-write Equation E.19 in terms of these values, in particular splitting the condition on $c_{t}$ being $j^{\prime}$-top for $j^{\prime}>j$ into two separate conditions,

Recall that h_{t,j^{\prime}}=\text{tanh}(\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})) indicates that $c_{t-1}$ is $j^{\prime}$-top, by the induction hypothesis. With our equation above we outline two different conditions that we’ll show together are necessary and sufficient to guarantee that the top slot is located at some $j^{\prime}>j$. Under each condition, we set the output gate to because of that (since if the top slot is at $j^{\prime}>j$, then slot $j$ is not equal to , and must not be let through the gate.)

While we cannot condition directly on the value of $c_{t}$, we know that $c_{t}$ is $j^{\prime}$-top for $j^{\prime}>j$ under two conditions on $h_{t-1}$, and thus on $c_{t-1}$. First, $c_{t-1}$ can be $(j+2)$-top or greater; since only one element can be popped at once, $c_{t}$ can be no less than $(j+1)$-top. Second, $c_{t-1}$ can be $j$-top or $(j+1)$-top, and the input $w_{t}$ is not some {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}; that is, it does not pop the top element off of the stack. Because $w_{t}$ is not {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, it must be an open bracket, {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, pushing to the stack. Hence, under these conditions, $c_{t}$ must be at least $(j+1)$-top. If neither of these conditions hold, then $c_{t}$ cannot be $j^{\prime}$-top for $j^{\prime}>j$, since if $c_{t-1}$ is $j$-top or $(j-1)$-top and w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, then $c_{t}$ must be at most $j$-top. And if $c_{t-1}$ is $j^{\prime}$-top for $j^{\prime}<j$, then it is impossible since only one bracket can be pushed at a time. Thus Equation E.27 is equivalent to Equation E.19, as required.

We now prove that parameters $(W_{o},U_{o},b_{o})$ implement Equation E.19, by implementing Equation E.27. In the first condition, h_{t-1,j^{\prime}}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}) for some $j^{\prime}>j+1$ and $i$; we want to show that $o_{t,j}=\mathbf{0}$. Based on our construction of parameters, we have that $o_{t,j}$ is upper-bounded by the following, when w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}:

where the last equality holds under our finite precision $\sigma$ by setting the scaling factor $\lambda$ such that $-0.5\lambda\gamma<-\beta$, that is $\lambda>\frac{2\beta}{\gamma}$.

In the next condition, we have that $h_{t-1,j^{\prime}}=e_{i}$for $j^{\prime}\in\{j,j+1\}$, and $i$, and x_{t}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}. In this case, we can write out the value of $o_{t,j}$ as follows:

Finally, if neither of those conditions hold, we have that h_{t-1,j^{\prime}}\not=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}) for any $j^{\prime}\geq j$, and so $h_{t-1,j^{\prime}}=\text{tanh}(\mathbf{0})=\mathbf{0}$. In this case, we can lower-bound the value of $o_{t,j}$:

as required. This completes the proof of the output gate.

The equation defining $f_{t,j}$ is as follows:

We define $W_{f,j,j^{\prime}}$ as follows:

again relying on the embedding-validity property.

that is, each row of $U_{f}$ is equal to the negative of the sum, over all close brackets, of the embedding of that close bracket (transposed.) Since all x_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}} are orthogonal and unit norm, this gives the following intermediate value:

that is, all units of the bias are equal to the same value.

We now prove that, as given, the parameters $(W_{f},U_{f},b_{f})$ implement Equation E.16 assuming the induction hypothesis of Lemma 4.

Equation E.16 has two cases. In the first case, $c_{t-1}$ is $j$-top and w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, and we need to prove $f_{t,j}=\mathbf{0}$. From the induction hypothesis, this means that h_{t-1,j}=\text{tanh}(\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})), and all other slots $h_{t-1,j^{\prime}}=\mathbf{0}$. The value of the gate is as follows:

where the last equality holds by setting $\lambda$ as already stated: $\lambda>\frac{2\beta}{\gamma}$.

The second case is whenever the conditions of the first case don’t hold, in which case we must prove $f_{t,j}=\mathbf{1}$. Thus, we either have $c_{t-1}$ not $j$-top, in which case,

or we have w_{t}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, in which case,

The equation defining $i_{t,j}$ is as follows:

We define $W_{i,j,j^{\prime}}$ as follows:

once again relying only on the encoding-validity property.

Since all x_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}} are orthogonal and unit norm, this gives the following intermediate value:

We now prove that the parameters $(W_{i},U_{i},b_{i})$, when plugged into $f_{\theta}$, implement Equation E.18 assuming the induction hypothesis of Lemma 4. Equation E.18 has three cases.

The condition of the first case is that $c_{t-1}=\mathbf{0}$, $j=1$, and w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}; we need to show that $i_{t,1}=\mathbf{1}$. Since $c_{t-1}=\mathbf{0}$, we have that $h_{t-1}=\mathbf{0}$. We can compute $i_{t,1}$ as follows:

note the bias term $-0.5\lambda\gamma$ is only set that way for $j=1$.

The condition of the second case is that $c_{t-1}$ is $(j-1)$-top, and w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}; we need to show that $i_{t,j}=\mathbf{1}$. From the induction hypothesis, we have that h_{t-1,j-1}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i^{\prime}}) for some $i^{\prime}$, and $h_{t-1}$ is elsewhere. We calculate $i_{t,j}$ as follows,

where we do not have to consider the case $i_{t,1}$ since $c_{t-1}$ cannot be -top, so this case cannot hold.

Finally, in the third case, none of the above conditions hold, and we need to prove $i_{t,j}=\mathbf{0}$. There are a number of possibilities here to enumerate. First, let $c_{t-1}=\mathbf{0}$ and $j=1$, but w_{t}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}. Then,

Next, we let $c_{t-1}=\mathbf{0}$ and w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, but $j>1$. Then,

Next, we let $j=1$, w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, but $c_{t-1}\not=\mathbf{0}$. Thus $c_{t-1}$ is $j^{\prime}$-top for some $j^{\prime}$.

If the second case of Equation E.18 doesn’t hold, then either $c_{t-1}$ is not $(j-1)$-top, or w_{t}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}; we need to prove $i_{t,j}=\mathbf{0}$. If $c_{t-1}$ is not $(j-1)$-top and $j>1$, then we can upper-bound the value of $i_{t,j}$ as follows:

If $c_{t-1}$ is $(j-1)$-top but w_{t}\not={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, then the value of $i_{t,j}$ is as follows,

Summary

We’ve now specified all parameters of the LSTM and completed the induction step in our proof of the stack correspondence lemma, Lemma 4; this completes the proof of the lemma.

Appendix F Proving generation in O​(m​k)𝑂𝑚𝑘O(mk) hidden units

For both the Simple RNN and the LSTM, we’ve proven stack correspondence lemmas, Lemmas 3, 4. These guarantee that for all prefixes of strings in Dyck-($k$,$m$), we can rely on properties of the hidden states of each model to be perfectly informative about the state of the DFA stack.

The two constructions differ in where they make information accessible; we describe those differences here and then provide a general proof that is agnostic to which construction is used.

Recall that we must show that the $\epsilon$-truncated support of $f_{\theta}$ (for both Simple RNN and LSTM) is the same set as Dyck-($k$,$m$). The token-level probability distribution conditioned on the history is given as follows:

and we denote this distribution $p_{f_{\theta}}$. We first specify $V$ by row; a single row exists for each of the $2k+1$ words in the vocabulary. We let $v_{w,j}$ refer to the row of word $w$, and the $k$ columns that will participate in the dot product with the rows of stack slot $j$ in $h_{t-1}$.

We start with the only difference in the softmax matrix between the Simple RNN and the LSTM. For the Simple RNN:

because the top of the stack is guaranteed to be in the first stack slot, and where $\zeta$ is a positive scaling constant we’ll define later. And for the LSTM,

for all $j$, since exactly one slot is guaranteed to be non-empty (if the stack is non-empty) and it could be at any of the $m$ slots. The rest of the softmax construction is common between the Simple RNN and the LSTM:

So that we can swap out $V$ when we swap out $\psi$ for a more efficient encoding, we state here a properties of $V$ we rely on (once $\psi$ is fixed):

where $u$ specifies the stack slot where the construction (Simple RNN or LSTM) stores the element at the top of the stack; for the Simple RNN, $u=1$, for the LSTM, $u$ is such that $c_{t}$ is $u$-top, as defined. This ensures that the softmax matrix correctly distinguishes between the symbol, {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, that is encoded in the top of the stack, from any other symbol. Further, we rely on

The softmax matrix we’ve provided obeys the softmax validity property with respect to $\psi$ since $\zeta e_{i}^{\top}e_{i}=\zeta$, and $\zeta e_{i}^{\top}e_{j}=0$, $j\not=i$. Further, $-\zeta\mathbf{1}^{\top}e_{i}=-\zeta$ for all $i$.

We’re now prepared to state the final lemma in our proof of Theorems 4, 4.

Let $w_{1:T}\in D_{m,k}$. For $t=1,...,t$, let $q_{t}=D_{m,k}(w_{1:T})$. Then for all $w\in\Sigma$,

Intuitively, Lemma 5 states that $f_{\theta}$ assigns greater than $\epsilon$ probability mass in context to tokens $w_{t}$ such that the prefix $w_{1:t+1}$ is the prefix of some member of $D_{m,k}$.

We proceed in cases by the state, $q$. We’ll show that a lower-bound on the probabilities of allowed symbols is greater than an upper-bound on the probabilities of disallowed symbols. We should note, however, that this is effectively a technicality to ensure our construction is fully constructive – that is, we provide concrete values for each parameter in the model; else, we could simply indicate that as $\zeta$ grows, the probability mass on $w$ such that $\delta(q,w)=r$ converges to , while the mass on all other symbols converges 1 over the number of such symbols.

Case q=[]𝑞q=[]

First, consider the case that $q_{t-1}=[]$. For all {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i})=r, and for all other symbols $w$, $\delta(q,w)\not=r$. By the stack correspondence lemmas, Lemma 4 3, we have that $\mathcal{Q}(h_{t-1})=q_{t-1}$ (Simple RNN) or $\mathcal{Q}(c_{t-1})=q_{t-1}$ (LSTM), and $h_{t-1}=\mathbf{0}$.

For all {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}, we have \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i})=r, and we have logits:

as required. For all w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, we have \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})\not=r, and we have logits:

Finally, for $\omega$, $\delta(q,\omega)\not=r$, and we have logits:

as required. With the logits specified for our whole vocabulary, we can compute the partition function of the softmax function by summing over all $k$ open brackets, all $k$ close brackets, and the END bracket, to determine probabilities for each token,

while the probability for any {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} and $\omega$ is p_{f_{\theta}}(w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}|h_{t-1})=e^{0.5\zeta\gamma}/Z_{[]}. The partition function is computed as $e^{0.5\zeta\gamma}$ multiplied by the number of words $w$ such that $\delta(q,w)\not=r$ (in this case, $k$ open brackets and the end bracket), plus a quantity upper bounded by $e^{-0.5\zeta\gamma}$ multiplied by the number of words $w$ such that $\delta(q,w)=r$ (in this case, $k$ close brackets.) All probabilities under this model are of this form.

Next, consider the case that q=[{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{1}},...,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i_{m^{\prime}}}] for $0<m^{\prime}<m$. For the Simple RNN, we have by Lemma 3 that h_{t-1,1}=\psi^{-1}({\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{m}^{\prime}) (the top of the stack.) Likewise for the LSTM, we have by Lemma 3, that $h_{t-1,m^{\prime}}=\text{tanh}(e_{i_{m^{\prime}}})$, and $h_{t-1,j\not=m^{\prime}}=\mathbf{0}$. Either way, we have v_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i}}^{\top}h_{t-1}\geq\zeta\gamma by the softmax validity property.In particular, for the Simple RNN, this expression is equal to $\zeta$ due to the softmax validity property; because of the tanh in $h_{t}=o_{t}\odot\text{tanh}(c_{t})$ in the LSTM expression, it’s equal to $\zeta\gamma<\zeta$ where $\gamma=\text{tanh}(1)$.

As with the first case, we construct the logits for each of the symbols. First, we have \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{m^{\prime}})\not=r; this is the top element of the stack, which can be popped with a close bracket. We have logits:

as required. For w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i} where $i\not=i_{m^{\prime}}$, we have \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})=r (since {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i} is not the top of the stack, seeing close bracket {\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\rangle}_{i} would break the well-balancing requirement), and the logits:

again by the softmax validity property. For w_{t}={\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}, we have \delta(q,{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i})\not=r, and the logits:

as required, since $m^{\prime}<m$, and by the softmax validity property, v_{{\color[rgb]{0,0.453125,0.8515625}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.453125,0.8515625}\langle}_{i}}^{\top}h_{t,m^{\prime}<m}=\mathbf{0}. Finally, for $\omega$, we have $\delta(q,\omega)=r$, and the logits