Stack Attention: Improving the Ability of Transformers to Model Hierarchical Patterns
Authors: Brian DuSell
David Chiang
Venue: ICLR 2024 spotlight (2024)
Keywords: transformer
attention
context-free languages
pushdown automata
formal languages
language modeling
machine translation
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose new attention operators that use stacks as a latent model of syntax, improving transformer language modeling on context-free languages and natural language.

Abstract: Attention, specifically scaled dot-product attention, has proven effective for natural language, but it does not have a mechanism for handling hierarchical patterns of arbitrary nesting depth, which limits its ability to recognize certain syntactic structures. To address this shortcoming, we propose stack attention: an attention operator that incorporates stacks, inspired by their theoretical connections to context-free languages (CFLs). We show that stack attention is analogous to standard attention, but with a latent model of syntax that requires no syntactic supervision. We propose two variants: one related to deterministic pushdown automata (PDAs) and one based on nondeterministic PDAs, which allows transformers to recognize arbitrary CFLs. We show that transformers with stack attention are very effective at learning CFLs that standard transformers struggle on, achieving strong results on a CFL with theoretically maximal parsing difficulty. We also show that stack attention is more effective at natural language modeling under a constrained parameter budget, and we include results on machine translation.

Published as a conference paper at ICLR 2024
STACK ATTENTION:
IMPROVING THE ABILITY OF TRANSFORMERS TO MODEL HIERARCHICAL PATTERNS
Brian DuSell"
David Chiang
Department of Computer Science ETH Zürich
Department of Computer Science and Engineering University of Notre Dame dcht ang@nd.
brian. dusel 1
1@ @iinf. et.z
ABSTRACT
Attention, specifically scaled dot-product attention, has proven effective for natu- ral language, but does not have a mechanism for handling hierarchical patterns of arbitrary nesting depth, which limits its ability recognize certain syntactic structures. To address this shortcoming, we propose stack attention: an atten- tion operator that incorporates stacks, inspired by their theoretical connections context-free languages (CFLs). We show that stack attention is analogous to stan- dard attention, but with latent model of syntax that requires no syntactic super- vision. We propose two variants: one related to deterministic pushdown automata (PDAs) and one based on nondeterministic PDAs, which allows transformers recognize arbitrary CFLs. We show that transformers with stack attention are very effective at learning CFLs that standard transformers struggle on, achieving strong results on a CFL with theoretically maximal parsing difficulty. We also show that stack attention more effective at natural language modeling under constrained parameter budget, and we include results on machine translation.
INTRODUCTION
Although transformers (Vaswani et al., 2017) have proven very successful on natural language, lin- guists have long held that language contains hierarchical syntactic structures (Chomsky, 1956, inter alia). Transformers do not appear to have any explicit mechanism for dealing with nested syntactic patterns of arbitrary depth. Recent thcoretical work, in fact, has shown that finite-precision trans- formers cannot recognize certain syntactie patterns, such Dyck-2 (Hahn, 2020), although they can do so up to bounded depth (Yao et al., 2021). Recent work has also shown that transformers have linear, rather than hierarchical, inductive bias (Petty & Frank, 2021: Mueller al., 2022) unless trained long past convergence (Murty et al., 2023a). This might explain, at least partly, why transformers are much less data-efficient than human children learning language (Gilkerson al., 2017: van Schijndel et al., 2019; Zhang et al., 2021; Frank, 2023; Warstadt et al., 2023).
In this paper, we propose novel type of attention mechanism, stack attention, that is explicitly designed to model hierarchical patterns. It does this by treating input vectors as items in a stack, and performing a soft-selection not directly over input vectors, but over sequences of stack ac- tions. We motivate this theoretically by pointing out that pushdown automata (PDAs), which are finite automata augmented with stacks, recognize the entire class of context-free languages (CFLs), which capture the essence of compositionality and recursion in natural language syntax (Hoperoft & Ullman, 1979; Sipser. 2013). Accordingly, we expect stack attention to increase transformers' expressive power, in addition possibly learning language from fewer examples and generalizing to held-out combinations of syntactic patterns in more human-like fashion.
Stack attention draws from prior work differentiable stacks, which are continuous functions that approximate the behavior of discrete stacks (Grefenstette et al., 2015; Joulin & Mikolov, 2015; DuSell & Chiang, 2020). We adapt two differentiable stacks into attention operators: the "super- position" stack of Joulin & Mikolov (2015), and nondeterninistic generalization introduced by
Work done while the University of Notre Dame.

---

Published as a conference paper at ICLR 2024
DuSell & Chiang (2023). Previous work attached differentiable stacks to recurrent neural networks (RNNs) as external memory modules, but here we integrate differentiable stacks into transformers in much deeper way. Whereas standard attention performs a soft-sclection differentiable stacks perform a soft-selection over sequences of push and pop actions on stacks of vectors, returning the topmost vector of the stack of the selected sequence. Accordingly, we differentiable stacks as drop-in replacement for standard attention, resulting in a transformer with latent model of syntax that, unlike most prior work, requires no syntactic supervision, and, in the case of nondeterministic stack attention, is able to recognize the entire class of CFLs.
imput vectors, these
Our experiments show that transformer language models with nondeterministic stack attention learn CFLs very effectively, consistently outperforming baseline transformers, and outperforming even the previously-pe pproposed VRNS-RNN (DuSell & Chiang. 2022) on CFI. with maximal parsing difficulty (Greibach, 1973). We also show that under constrained parameter budget, transformer language models with stack attention outperform baseline transformers on the Penn Treebank by 4.3 perplexity points, and we include results on small machine translation task based on the German- English dataset from Europarl v7. Our code is publicly available."
2 RELATED WORK
Our method enjoys two key advantages over prior work on syntax-oriente transformers: (1) it does not require syntactic supervision (although can be extended leverage it); and (2) it is unidirec- tional in time, so it can be used in causal language models and decoders. Most prior work requires syntactic supervision (Shiv & Quirk, 2019; Deguchi et al., 2019; Zhang et al., 2020; McDonald & Chiang, 2021; Qian et al., 2021; Sartran et al., 2022: Murty et al., 2023b). Wang et al. (2019b) proposed an unsupervised but bidirectional model. Kim et al. (2017) proposed a framework for structured attention mechanisms that sum to our approach. Their work included an attention operator based on the inside-outside algorithm for projective dependency trees, whereas rithm for PDAs. A key difference is that, at each timestep, their attention mechanism marginalizes over all latent structures backward and forward in time, so it bidirectional. Our stack attention marginalizes backward in time only.
exponential number of latent structures, similar
propose one based on dynamic programming algo-
3 BACKGROUND
We start by reviewing attention and differentiable stacks. For any tensor (matrix, vector) A, let Ali, j....] indicate the element or sub-tensor indexed by i.J..... Let ar be the logistic function.
3.1 SCALED DOT-PRODUCT ATTENTION
Attention has become a cornerstone of modern neural network architectures (Cho et al., 2015; Bah- danau et al., 2015; Vaswani al., 2017). Loosely defined, attention refers an operation that, given a sequence of values v Vn E Rdd,, produces an output v' that represents a linear interpolation, or "soft-selection." over Ui..... Uy..
The transformer uses an attention operator called scaled dot-product attention (SDPA). The self- attention sublayer function in a transformer receives input vectors 2' , x²,, E Ddmaddl and linearly transforms them to sequences of query vectors ....... 9r. E dR key vectors k1..... kn E Rd and value vectors V1.... ... Un E RR., It runs SDPA for each timestep < t n, producing outputs vf...., 2n, E Rdd. Each 2t is linearly transformed to an output y/ Paoddd
qi Wa
k; Wi
11
W.
(1)
q: k;
<1
2: = softmax(21)
(2)
Vdl
2 [i]
v Wy
(3)
2

---

Published as a conference paper at ICLR 2024
If prevent SDPA from attending to future timesteps by ending the iterations over at instead of n, which is necessary for causal language modeling and decoding, we say is causally masked. Each sublayer typically has multiple attention heads, applying multiple SDPA operators to the input sequence and adding the results together.
3.2 DIFFERENTIABLE STACKS
stack is container of elements that allows one push and pop elements in fat-in-first-out order. A differentiable stack is a continuous function that upproximates the behavior of a stack: in this paper, we assume elements are vectors in R Suppose we start with an empty stack So and perform operations on it (e.g. pushing and popping), resulting in sequence of stacks , In a differentiable stack, we represent the actions as stack action vectors a1,.... an where each a, is a set of action weights. The exact set of available actions varies with the particular style of differentiable stack. The inputs also include a sequence pushed vectors vi...., Un E R that serve the new elements inserted by push actions. The output of the differentiable stack is sequence stack reading vectors T1. Tn, where each Te represents the top vector element of the stack after processing the stack actions and pushed vectors up to timestep l. Each differentiable stack can be abstracted into a function STACK that incrementally updates the stack, and a function READING that queries the stack reading from it.
S. = STACK at,
T: = READING(S.)
(4)
The stack readings differentiable with respect to the stack actions and pushed vectors, if the stack is incorporated into a neural network, the whole network can still be trained end-to-end with backpropagation, without any need for supervision over stack actions.
3.2. SUPERPOSITION STACK
The differentiable stack of Joulin & Mikolov (2015) uses strategy called superposition. It sim- ulates sittiinnall-weiggteee stack actions by computing three new, separate stacks: one with all elements shifted down (push), kept the same (no-op), and shifted up (pop). The new stack is an element-wise interpolation ("superposition"" "" of these three stacks. Another way of viewing this is that each element is an interpolation of the elements above, and below it at the previous timestep, weighted by push, no-op, and pop probabilities respectively. Here, a, [a.ush a; a. aP"PjT is probability distribution over the three actions, and the stack state S, is matrix V, containing the stack elements. We implement V2 STACK (V-1,, a2, Vt) as
poop
Vi[il = at push (BOVE,(1) + di
noop (Tx Pop BELOW,(1) (15i50)
(5)
i=1
(B)
(6)
V Vi-1ii - 1] i> 1
V [1] i<t
(T
0
i=t
Vi li
+
1]
i<t-1
BELOW
(8)
i2t-1.
The stack reading r, READING V,) is simply the top vector element V,[1], and we V, [0]. This stack has quadratic time and space complexity with respect to input length.
3.2.2 NONDETERMINISTIC STACK
DuSell & Chiang (2023) recently proposed differentiable stack based on pushdown automata (PDAs). PDAs are type of nondeterministic automaton that consists of finite state machine connected to an infinite stack memory (for a detailed definition, see Appendix A). They recognize exactly the class of CFLs. Every PDA has finite alphabet of input symbols Σ, finite set of states Q, and a finite alphabet of stack symbols The initial stack consists of 1 E Т, which is a designated bottom symbol.
This presented little differently from Vaswani et al. (2017), but it mathematically equivalent.
3

---

Published as a conference paper at ICLR 2024
A PDA has transitions of the form q, S r, e, which signifies that if the PDA is in state g Q and has string u E I on top of the stack, and string a E EU{E} can be scanned from the input tape, then it goes to state in the sense that multiple transitions may apply ambiguously to the same machine configuration, in which case all sequences of transitions are explored. Nondeterminism is essential for recognizing all CFLs, because deterministic PDAs are strictly less powerful. We call a sequence of connecting transitions that starts from the initial machine configuration run. We write ~r i, q, x to indicate that run π scans exactly t input symbols, ends in state 9, and ends with symbol x on top of the stack.
Q, pops u, and pushes string E - onto the stack. PDAs are nondeterministic
The differentiable stack of DuSell Chiang (2023) assumes the PDA is in the following normal form, which is equivalent in recognition power. Transitions must have one of the following forms, where q, r E Q a E Σ, and r, y T::
q, I > r: xy
Push y on top of x.
q. S V. V
Replace z with Pop x.
(10)
q, x 4 T,E
(11)
Weighted PDAS (WPDAs) are an extension of PDAs that assigns a non-negative weight to each transition. The weight of run π,, which we denote (T)),, is the product of its transition weights. DuSell & Chiang (2023) extended WPDAs further by augmenting every instance of a stack symbol with a vector in RT so that stack elements are members of I x R". The purpose of adding vectors is to allow the automaton to transmit information via the stack efficiently using embedding vectors rather than a finite alphabet of discrete symbols. We will refer to as the vector PDA (VPDA).
Let w = W1 Wn E E be a string of input symbols, and assume that there is a sequence of pushed vectors vo.... Un. In a vector PDA, the initial stack consists of vo). Transitions have the following semantics when scanning Wt:
q,
Wt r, xy
If (z, u)) is on top, push (y, v,).
r,
If (ar, u)) is on top, replace it with (y, u.)
q T,
If (2r, u)) is on top, pop it.
Note that transitions can only be conditioned on (q.23 E Q x r, never the value of vector u, so only Q and l direct the nondeterministic branching of the VPDA (this will allow Eq. (12) to be tractable).
The differentiable stack of DuSell & Chiang (2023) is a differentiable simulation of a vector PDA. We will refer it as the differentiable vector PDA (dVPDA). Here, each a, is the flattening of tensor A, which contains the transition weights used when scanning Wy, We denote the weight of
transition q, x - r, A Ag[g, + v). Since g, r E Q, E T, and there are 2 2 + possibilities
for e given Eqs. (9) to (11), the size of A, is Q x T x [Q] x (2[T) + 1). We set ey = ouw,)
where w. E I is learned parameter.
Let uTT) denote the top stack vector at the end of VPDA run ग. The stack reading E 1Q includes, for each (r, y) E Q I, an interpolation of uTT) for every run π ~+ t, r. y, normalized by the weight of all runs. Let rz (r, u] denote the slice of r. corresponding to (r,y)- Then
Exetery (bT)) v(x)
r relr. v] EneEQ Ever Exner' (TT)
(12)
VPDA runs that outputs the expected top stack vector. Although Eq. (12) sums over an exponential number of runs, it can be computed in cubic time and quadratic space using Lang's dynamic programming algorithm (Lang, 1974), which can be expressed using the abstraction of Eq. (4), See Appendix B for details.
In this way, the dVPDA is soft-selection
4 COMPARISON OF DIFFERENTIABLE STACKS
Here we make the novel insight that the superposition stack is a special case of the dVPDA. If unroll the superposition stack's equation for Tı, we see that it is summation that enumerates all possible sequences of actions:
r = Σ (T) u(T)
(13)
Two
4

---

Published as a conference paper at ICLR 2024
where is of VPDA with [Q] ||| 1, and iT ~e t means that ends at timestep t. This means the superposition stack is a special case of the dVPDA with [Q] = |T] = 1. normalized transition weights, and vn = 0. So, like the dVPDA. the superposition stack a soft-selection over VPDA runs, but for a weak VPDA where nondeterministic branching can only be directed by the choice to push, pop, do nothing each step, and not by the state machine stack contents.
This explains differences in the range of tasks these two stacks can solve, which we will observe in Section 6. Both stacks are real-time, meaning that all PDA transitions scan exactly one input symbol. The superposition stack resembles real-time deterministic PDAS (DPDAs), which recog- nize only a subset of CFLs called real-time deterministic CFLs (DCFLs) (Ginsburg & Greibach, 1966: Igarashi, 1985). In contrast, nondeterministic PDAs recognize all CFLs even when they are real-time (Greibach, 1965) and in the normal form of Eqs. (9) to (11) (DuSell Chiang, 2023). Therefore dVPDAs, and transformers that contain them, can recognize all CFLss.
5
METHOD
We construct a stack attention layer by replacing multi-head SDPA with a differentiable stack. Trans- formers consist of multiple layers, each of which contains multiple sublayers. Every sublayer in- cludes layer normalization, dropout, and residual connections. In a transformer encoder, each layer consists of self-attention sublayer followed by a feedforward sublayer. Similarly, in a transformer decoder, each layer consists of a self-attention sublayer, a cross-attention sublayer that attends to the output of an encoder, and a feedforward sublayer, in that order. Let o Rdmdded and y: E modde be the input and output, respectively, of a sublayer at timestep t. A sublayer is implemented as
x' LAYERNORM(Z:)
(14)
y SUBLAYER(I, (x',....,,,, xr,)
(15)
yt 22, + DROPOUT(Y().
(16)
Here, SUBLAYER is called the sublayer function, and it customizes the behavior of the sublayer. In standard SDPA layer, the sublayer function is multi-head SDPA (Section 3.1). LAYERNORM is layer normalization (Ba et al., 2016). We use pre-norm instead of post-norm (Wang et al., 2019a; Nguyen& & Salazar, 2019). Layer normalization is also applied the output of the last layer.
Like SDPA, the superposition stack and dVPDA produce a weighted sum of input vectors (Sec- tion 4). Based on this insight, we adapt each differentiable stack into an attention sublayer by using it as the sublayer function instead of multi-head SDPA. We linearly transform each sublayer function input 2 E IR model to v. E I using the logistic function to ensure elements are in 10.
8. = (WW22
(17)
This is similar to the way SDPA generates values (ef. Eq. (1)), and like SDPA, stack attention performs a soft-selection over them. We also linearly transform xt to E Rd,, analogously to the way SDPA generates queries (cf. Eq. (1) again). For the superposition stack, we use softmax to ensure a, is probability distribution.
a softmaxxX
(18)
For the dVPDA, we ensure weights are non-negative using exp. (In implementation, we actually calculate run weights in log space to avoid overflow, and the exp is merely implicit.)
a: = exp Wax²)
(19)
We compute rn from a1,..., ag and U1...., v. according to Eq. (4), using STACK and READING for one of the two differentiable stacks. We call stack attention based on the superposition stack superposition stack attention, and we call stack attention based on the dVPDA nondeterministic stack attention. We linearly transform each r, E Rd. to get the sublayer function output yí E Rddoddd (cf. Eq. (3)).
(20) Note that because the differentiable stack does not attend to future timesteps, we do not need to apply any masking to enforce causality. We illustrate our stack attention architecture in Fig. 1. To characterize the practical speed of stack attention on GPUs, we show its parallel time complexity in Table 1. See Appendix C for more details time and space complexity.
y. = Wrr:
5

---

Published as a conference paper at ICLR 2024
output
Mt-1 DROPOUT
31
Mt+1
PT
DROPOUT T
DROPOUT
7
stack
...
St
S
S
: ac-1. 1--1
? a 4++1
LAveRNORM
LAVERNORM
LAVERNORM
1
input
21-1
4
21+1
Figure 1: Conceptual diagram of stack attention sublayer, unrolled across portion of time. Dotted arrows indicate linear transformations, and dashed arrows indicate residual connections.
Table I: Parallel time complexity of the types of attention studied this paper, as a function of sequence length n. "Implemented" shows the complexity of each implementation used in this paper. Nondeterministic stack attention could be further parallelized a manner similar to parallel CKY (Kosaraju, 1975), which is impossible in RNNs ("Parallel CKY)). Since CFL recognition is in NC (Ruzzo, 1981), theoretically this could be lowered even further to O(log ("Theoretical").
Attention
Implemented
Parallel CKY
Theoretical
SDPA
O(log
Superposition
O(n)
0 ((logn) )))
Nondeterministic
O
(nn log n)
((logn)
6 CONTEXT-FREE LANGUAGES
In this section, we test the performance transformers with stack attention as language models on the same five CFL tasks from DuSell & Chiang (2020; 2022). Let wt denote the reverse of w.
w ## R The language uw | w E {0,1}*} www R The language (ww" w E {3"1
w ft Like ww but with higher tendency to have long stretch of the same symbol repeated in the middle. Strings are the form waPw where w E {0, 1,",,,,"00 """""aa a E {0, and p 2 0.
Dyck The language of strings with two types of balanced brackets.
Hardest CFL A highly ambiguous, CFL-complete language with maximal parsing difficulty (Greibach, 1973). See DuSell & Chiang (2020) for details.
The w R and Dyck languages are real-time DCFLs, whereas ww waww and Hardest CFL are nondeterministic CFLs. We use the same sampling procedure as DuSell & Chiang (2020) to generate datasets for each task. Every time we train a model, randomly sample training set of 10k examples and validation set of Ik examples, both with lengths in the range [40, 80). For each task, we sample test set with string lengths varying from 40 to 100. with 100 examples per length. The test sets remain the same across all runs. Following DuSell & Chiang (2022), we evaluate models using the cross-entropy difference between the distribution the model learns and the true distribution the data is sampled from. Units are nats, lower is better, and 0 is optimal.
We compare transformers with SDPA (Tf), superposition stack attention (Tf+Sup). and nondeter- ministic stack attention (Tf+Nd), as well as their stack-augmented L.TTM (Hochreiter & Schmid- huber, 1997) counterparts (LSTM&NN what DuSell Chiang (2023) call the VRNS-RNN). All transformers have 5 layers. For transformers with stack attention, in the third (middle) layer, replace the SDPA sublayer with the corresponding stack attention sublayer. Although we could have used stack attention in all layers, only is necessary for recognizing CFLs, and multiple stack attention layers would be computationally costly. Using stack attention in the middle layer is compromise between placing it at the beginning end of the transformer.
We have adjusted model sizes so that their parameter counts satisfy the following constraints on the Hardest CFL: (1) each type of stack has fewer parameters than the preceding types (none Sup >

---

Published as a conference paper at ICLR 2024
Table 2: Results of training transformers with different attention mechanisms as language models on the Penn Treebank benchmark of Dyer et al. (2016), measured with perplexity. All results are the best of 20 random restarts, selected by validation cross-entropy. Despite having the fewest parameters, our nondeterministic stack attention (Tf+Nd) has the lowest perplexity.
Model
Params.
Val. 4
Test +
Tf
10,051 ,072
115.11
92.84
Tf+Sup (Ours)
10,050,304
122.94
98.67
Tf+Nd (Ours)
9,861,898
110.59
88.54
Nd); and (2) each transformer has fewer parameters than its LSTM counterpart. For transformers, the feedforward hidden layer size is 2 dindeel SDPA layers causally masked and have 4 heads. For Tf, we use dmodel = 32. For Tf+Sup. we use dmodel = m = 32. For Tf+Nd, we use dmodel = 28, m = 5, and the same values for [Q] and [I] as DuSell & Chiang (2022) ([Q 2, [[| 3 for w#w ww. and Dyck; [Q] = 3. |[] = 3 for waPw and Hardest CFL). For LSTM, we use one layer with 100 hidden units. For Sup, we 93 hidden units and m 10. For LSTM+Nd,
we use 64 hidden units and the same dVPDA sizes as Tf+Nd. See Appendix D for more details.
For each language and architecture, we train 10 models and report results for the model with the best validation performance (Fig. 2). Tf+Sup outperforms Tf on all tasks except Hardest CFL. TY+Nd achieves strong results across all tasks, except of-distribution lengths on DCFLs (w#w and Dyck). Notably, Tf+Nd outperforms all models on Hardest CFL for in in-distribution lengths despite having the fewest parameters. It also outperforms Tf and Tf+Sup on all nondeterministic CFLs (ww.. woPw Hardest CFL). Although multi-head SDPA can express multiple interpretations of an input at once, represents only a fixed number, whereas Tf+Nd sums over an exponential number of VPDA runs. The poor performance of Tf and even Tf+Sup ww"" an extremely simple nondeterministic task, highlights the importance of nondeterminism.
How do transformers compare to their LSTM counterparts? For all tasks except Dyck, LSTM out- performs Tf, and transformers have serious length generalization issues. This may have to do with the linear bias in the SDPA layers, and different positional encodings might also alleviate this (we use sinusoidal encodings). However, Tf+Nd appears to alleviate this problem on nondeterministic CFLs, but it still overfits the training length distribution on the DCFLs. LSTM+Sup generally outperforms Tf+Sup. LSTM+Nd has better length generalization than Tf+Nd and outperforms on ww"", but Tf+Nd has better in-distribution performance on Hardest CFL.
Qualitatively, we find that Tf+Sup learns easily interpretable action patterns on the real-time DCFLs. On w R2 it learns to push before & and pop afterwards. On Dyck, it learns to push after reading opening brackets and pop after reading closing brackets. See Appendix E for details.
NATURAL LANGUAGE MODELING
In this section, we test transformers with and without stack attention on a natural language modeling benchmark. We use the Penn Treebank (Marcus et al., 1994) as preprocessed by Dyer et al. (2016). We train the models to learn a probability distribution over single sentences, without context from past sentences. Each transformer has 5 encoder layers, with stack attention inserted in the same way as Section 6. We use dmoodel 256 and a feedforward hidden layer size of 1024. SDPA sublayers are causally masked and have 8 heads. For superposition stack attention (Tf+Sup), we set m 511. For nondeterministic stack attention, we set [Q] = 3, [[] = 3, and m = 10. These settings ensure that the number of parameters in each model is less than that of the preceding baselines. Note that these models are relatively small; for comparison, the Transformer-X XL language model of Sartran et al. (2022) has 12M parameters. See Appendix F for more details, including computational cost.
For each architecture, we train 20 models and report results for the model with the lowest valida- tion cross-entropy (Table 2). We see that Tu+ does not outperform the Tf baseline, but Tf+Nd outperforms both, despite having the fewest parameters. This suggests that nondeterministic stack attention is more efficient encoding language under constrained parameter budget.
7

---

Published as a conference paper at ICLR 2024
LSTM
Sup Tf+Sup (Ours)
LSTM+Nd
Tf
Tf+Nd (Ours)
Cross-entropy Dift.
0.4
Cross-entropy Diff.
2.0
w
0.3
1.5
0.1
0.5
0
20
40
60
80
100
50
60
70
80
90
100
Cross- Cross-entroppe eate-eetrop Diff.
Cross-ene cross-entropy Diff.
0.8
wig
0.3
0.4
0.0
0.0
100
150
50
70
90
100
Cross-entrop Diff.
waP
Cross Cross-entropy Diff.
w00 un
0.6
0.4
0.1
0.2
0.0
0.0
0
50
100
150
40
50
60
70
80
90
100
Dyck
Dyck
Cross entsoyy Diff.
Cross-cntropy Diff.
0.1
0.0
80
50
70
90
100
Cross- entropy Diff.
Hardest CFI.
Cross-cntie cntropy Diff.
0.3
Hardest CFL
0.1
0.1
100
200
50
70
Epoch
Length
Figure 2: Language modeling results on context-free languages, comparing transformers with and without stack attention, as well their LSTM counterparts. Left: Cross-entropy difference (L) in nats between model and source distribution on the validation set, as a function of training time. Lines the best of 10 runs, selected by validation cross-entropy difference. Right: Cross-entropy difference (L) on the test set, binned by string length. The dashed line indicates the longest length in the training set. See Table for model parameter counts. Nondeterministic stack attention (Tf+Nd) outperforms standard attention (Tf) on w R and Hardest CFL; and it achieves the best in-distribution performance Hardest CFL despite having the fewest parameters.
8

---

Published as a conference paper at ICLR 2024
Table 3: Machine translation results for models trained subset of 100k examples from the Europarl v7 de-en corpus. Models are organized into three tiers by number of parameters. All results are the best of 5 runs, selected by decoder cross-entropy on the validation set. We report decoder perplexity on the validation set (newstest2016) and BLEU on the test set (newstest2017).
Model Tf
dmodel
Params.
Val. Perp. +
Test BLEU
160
4,595,200
12.52
12.21
TF+Sup (Ours)
160
4,492,480
11.94
12.03
Tf+Nd (Ours)
160
4,465,610
13.00
11.86
Tf
240
9,580,800
12.54
12.11
Tf+Sup (Ours)
240
9,349 ,200
11.53
12.81
Tf+Nd (Ours)
240
9,232,810
12.46
11,50
360
20,419,200
11.80
12.69
TE+Sup (Ours)
360
19, 19,900,080
12.03
12.03
Tf+Nd (Ours)
360
19,551,610
12.54
11.74
MACHINE TRANSLATION
In this section, we test transformers with stack attention on a small machine translation task, using stack attention in both the encoder and decoder. We use subset of the German-Enge English dataset from Europarl v7 (Koehn, 2005), simulating a low-resource scenario. We use the news translation tasks from WMT 16 (newstest2016) and WMT 17 (newstest2017) 7) validation and test data, respectively. To accommodate the computational cost of nondeterministic stack attention, we limit the training, validation, and test
examples where both the source and target side have no more than 150 characters, and we randomly subsample 100k examples from the training set. We tokenize the data into a vocabulary of 32k tokens using BPE (Sennrich et al., 2016).
We use layers in both the encoder and decoder, swapping out SDPA with the same kind of stack attention in the third layer of each. We continue use SDPA the cross-attention mechanism. In order to examine model performance as function of parameter count, we vary dmodel to test models with three size tiers, roughly doubling the number of parameters in each successive tier, and ensuring that the number of parameters in each model is less than that of the preceding baselines. We use heads each SDPA sublayer and set the feedforward hidden layer size to dmodel- For superposition stack attention (Tf+Sup), we set (Tf+Nd), we set |Q| = 3, [I| = 3, and m 5. See Appendix G for more details.
dmodel. For nondeterministic stack attention
For each tier and architecture, we train 5 models and report results for the model with the lowest decoder cross-entropy on the validation set (Table 3). Stack attention does not appear to improve translation quality over Tf, although Tf+Sup performs best on the middle tier. Why does stack attention not help consistently? Perhaps natural language does not actually contain many deeply- nested hierarchies (especially syntax does not provide a benefit over the baseline transformer, as suggested by Yao et al. (2021). However, Section seems to indicate that this is not the case, least for language modeling. Perhaps language is hicrarchical, but baseline transformers of this size are already able to learn syntax just as effectively as those with stack attention. Or, perhaps stack attention can provide advantage over baseline transformers, but the training procedure simply does not find parameter settings that do so.
sentences limited 150 characters). so adding latent model of
CONCLUSION
We showed that two types of differentiable stack, the superposition stack (Joulin & Mikolov, 2015) and nondeterministic generalization of it (DuSell & Chiang, 2023), can be incorporated into trans- formers as attention mechanism. On CFL language modeling tasks, we showed that nondeter- ministic stack attention improves upon the standard transformer architecture, and even improves upon its RNN counterpart on the Hardest CFL, a challenging ambiguous language. We also showed that nondeterministic stack attention improves perplexity on a language modeling benchmark. We believe stack attention is important development in unsupervised syntax-oriented transformers.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
To facilitate reproducibility, we have publicly released all code we used to download and preprocess datasets, run experiments, and generate the figures and tables in this paper. During both develop- ment and experimentation, we ran our code containers to simplify the replication of our software environment. Our code includes the original Docker image definition commands we used for each experiment, figure, and table. We have thoroughly documented our experimental settings in Sections to 8 and Appendices D, F and G.
used, as well as the shell
Although we cannot distribute the Penn Treebank dataset used in Section 7, we have included a script that generates the same preprocessed files used by Dyer et al. (2016) and Sartran et al. (2022) from the raw Penn Treebank distribution files, so that anyone with a license for the Penn Treebank can reproduce them. To our knowledge, such a script was available previously.
We used publicly available datasets for the machine translation experiments in Section 8.
ACKNOWLEDGEMENTS
We thank Laurent Sartran for providing with the preprocessed Penn Treebank dataset used by Dyer et al. (2016) and Sartran et al. (2022), and for answering our questions. We thank Ken Sible for helpful discussion about the implementation of our machine translation system. We thank Darcey Riley, Aarohi Srivastava, Stephen Bothwell, and Ken Sible for their feedback on draft of this paper. We thank Alex Warstadt and Chihiro Taguchi for pointing to some of the cited work. We thank the Center for Research Computing at the University of Notre Dame for providing the computing infrastructure used for our experiments. Finally, we thank the anonymous reviewers for their invaluable discussion and feedback. This material is partially based upon work supported by National Science Foundation under Grant No. CCF-2019291.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. In NIPS 2016 Deep Learning Symposium, December 2016. URL https:// wwt.:/rrrttttttttt/// forum?
id=BJLa_2CC
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Proc. ICLR, San Diego, California, USA, May 2015. URL https:/ /tt.a/awwrgg/aaaaaaaa/aa/ 0473..473..
Kyunghyun Cho, Aaron Courville, and Yoshua Bengio. Describing multimedia content using attention-based encoder-decoder networks. Trans. Multimedia, 17(11):875-1886, Novem- ber 2015. doi: 10.1109//MMM.20552447000 URL document /7243333.
ieee.ors
Noam Chomsky. Three models for the description of language. IRE Trans. on Information Theory, 2(3):113-124, 1956. doi:
1956. 1056813. URL htlLtttss///eeeet// ieee.
org//ocument/ /1056813.
Hiroyuki Deguchi, Akihiro Tamura, and Takashi Ninomiya. Dependency-based self-attention for transformer NMT. Proc. RANLP, pp. 239-246, Varna, Bulgaria, September 2019. IN- COMA Ltd. doi: (02261159978-954445--666 019-1028.
URL https:
URLstts::/aclatthe httpss///caaatthoooooooo org/
Brian DuSell and David Chiang. Learning context-free languages with nondeterministic stack RNNs. In Proc. CoNLL, pp. 507-519, Online, November 2020. Association for Computa- tional Linguistics. doi: (0.18653/v1/2020.co 2020. co011-1.11
URL https: httt::pps//alllahhoooogyy org/
Brian DuSell and David Chiang. Learning hierarchical structures with differentiable nondetermin- istic stacks. In Proc. ICLR, Online, April 2022. URL https: /pwpnrevii.e...ceooo net forum? i id==5txw_@p1133
10

---

Published as a conference paper at ICLR 2024
Brian DuSell and David Chiang. The surprising computational power of nondeterministic stack RNNs. Proc. ICLR, Kigali, Rwanda, May 2023. URL
Forum?2d=o58JtcDeS.
Chris Dyer, Adhiguna Kuncoro, Miguel Ballesteros, and Noah A. Smith. Recurrent neural net- work grammars. In Proc. NAACL HLT, pp. 199-209, San Diego, California, USA, June 2016. Association for Computational Linguistics. doi: 10.1863//1/N11-102224 URL https: / //aclannthlooyy wttpr/hoooooyyoog/1----100000 ozq/N16-1024.
Michael C. Frank. Bridging the data gap between children and large language mod- els. Trends in Cognitive Sciences, 27(11):990 November 2023. doi: 10.1016/j.tics 2023.08.007. URL https: http:p/:/wwwwscceeeeeeeeeeeeeeee 91364661323002036
Jill Gilkerson, Jeffrey A. Richards, Steven F. Warren, Judith K. Montgomery, Charles R. Green- wood, D. Kimbrough Oller, John H. L. Hansen, and Terrance D. Paul. Mapping the early language environment using all-day recordings and automated analysis. American Journal of Speech-Language Pathology, 26:248 265, May 2017. doi: 10. 10.1044/2016 AISLP-15-0166 URL https::/ppuil asha .rg. doo /10. 044/ /244///1115 22t___11111111--111112111_1 1-
Seymour Ginsburg and Sheila Greibach. Deterministic context free languages. Information and Control, 96)):20-648, October 1966. doi:
URL https:
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proc. AISTATS, volume 9 of PMLR, pp. 249-256 Sardinia, Italy, May 2010. URL attps://proceedingg mlr press http:////gg///oooeeeeeee html.
Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learning to transduce with unbounded memory. In Advances in NIPS, volume 28, Montreal, Canada, De- cember 2015, Curran Associates, Inc. URL httttpss//paaeesssss http://aappersaaeee mntpersss.ppie. nips. cc/ cttp:perepppppeeeeeeeeeeeeee files/ Absttatt .ht.la.ttt.tm......m
Sheila A. Greibach.
normal-form theorem for context-free phrase structure grammars. J. ACM, 12(1):42-52, January 1965. doi: 10.1145/3212503322544,, URL https: //d1.acm. aacm brg/doi/10.. 1..0..144 1et.e1111122255 321254.
Sheila A. Greibach. The hardest context-free language. SIAM Computing, (4):304-310 1973.
doi: 10.1137/0202025. URL https: /epubs.....//////////ee siam. org/ /woo/oooo////////// 1137 1111/12020222222 0202022..
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Trans. ACL, 8:
156-171, Januarr 2020. doi: 10.1162/tacl. a.00306. URL https: /tttps://dci.oorg1000
ttcc_a_00306
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, (8):1735-1780, November 1997. doi: 111162/ncco.1997.99 1735. URL https: //ww.. researchgate... www.t//wwwcaaaio////o hare:a224-Looggooohttttteeee Sennttttrr_MMMorry... Memory.
John Hopcroft and Jeffrey D. Ullman. Introduction to Automata Theory, Languages and Compu- tation. KidisonnWesleye 1979.
Yoshihide Igarashi. A pumping lemma for real-time deterministic context-free languages. Theoret- ical Computer Science, 36:89-97, 1985. doi: 10.101609004499555339000333333 URL https:
/www.. sci sciencedi rect
Armand Joulin and Tomáš Mikolov. Inferring algorithmic patterns with stack-augmented re- current nets. In Advances in NIPS, volume 28, Montreal, Canada, December 2015, Cur- ran Associates, Inc. URL https: httttps::prroceeediiiie www.nn......iiiieeeeeeeeeee.ee.e F 9tw.befffe..7669999--99-6 Absttttt......a................. html.
neurips ips..ps..... wwt.ripp.c..pppprr//////aaa wwtp//paee/211////e
Yoon Kim, Carl Denton, Luong Hoang, and Alexander M. Rush. Structured attention networks. In Proc. ICLR, Toulon, France, April 2017. URL https: http:////ppeeeeee.e/mm net htt.:foruttftttmmee/// forumm HkEONvqlg.
11

---

Published as a conference paper at ICLR 2024
Diederik P. Kingma and Jimmy Lei Ba. Adam: A method for stochastic optimization. In Proc. ICLR, San Diego, California, USA, May 2015. URL https: /ttp:://aaxvvvvv. /atpp::/aarxl..oogg/aabssss/ /412. 6980.
Philipp Koehn. Europarl: parallel corpus for statistical machine translation. In Proc. MT Sum- mit, pp. 79-86, Phuket, Thailand, September 2005. URL https::/accaattaeeeetee
.org/
2005. . mt a summit committ-papeee wtt-apppppss....
S. Rao Kosaraju. Speed recognition of context-free languages by array automata. SIAM J. Com- puting, 4(3):331-344 340, 1975. doi: 10.1137/0204028. URL htt:p:://epule s iamm org/
doi/10 1137 2t1.11.1117000002222220 0204028.
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proc. EMNLP: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi:
(0.18653/v1/D18--1
URL httptthttcccaaatae htrtpp::/aclannthooogyy.og orggh/oooooooooooooooo 2012.
Bernard Lang. Deterministic techniques for efficient non-deterministic parsers. In Proc. ICALP, pp. 255-269, Saarbrücken, Germany, July 1974. Springer. doi: 10.00077998-33--22215555666
262--6555---222666-2- URL attps://lint springez. com wwtp://caaateeer111000////111/7/ 18.
Mitchell Marcus, Grace Kim, Mary Ann Marcinkiewicz. Robert MacIntyre, Ann Bies, Mark Fer- guson, Karen Katz, and Britta Schasberger. The Penn Treebank: Annotating predicate argu- ment structure. In Human Language Technology: Proceedings of a Workshop held at Plains- boro, New Jersey, March 8-11, 1994, Plainsboro, New Jersey, USA, March 1994. URL https://aclanthonn
org/ / 10200
Colin McDonald and David Chiang- Syntax-based attention masking for neural machine trans- lation. In Proc. NAACL: Student Research Workshop, pp. 47-52, Online, June 2021. Asso- ciation for Computational Linguistics. doi: (0.18653/112021..macck / //aclanthology. , wrg/ cxg/2021.naacl-srcc naacc-ssw..
URL https:
Aaron Mueller, Robert Frank, Tal Linzen, Luheng Wang. and Sebastian Schuster. Coloring the blank slate: Pre-training imparts hierarchical inductive bias to sequence-to-sequenee models. In Findings of ACL, pp. 1352-1368, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 00.18833/v1/002.. /206://////222..nnnnn---aa 106. URL https: http::sss///alaaahlloooooooo org/ 2022. findings-aannneeeeen acl, 106.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. Grokking of hierarchi- cal structure in vanilla transformers. In Proc. ACL (Short Papers), pp. 439-448, Toronto, Canada, July 2023a. Association for Computational Linguistics. doi: Eettss:/aclanttin
acloo.. URL
org/2023. acl short.38.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. Pushdown layers: Encoding recursive structure in transformer language models. In Proc. EMNLP, pp. 3233-3247 Singapore, December 2023b. Association for Computational Linguistics. doi: 10.18653/11/2223 emalp-main. 195. URL https:ttp.... /acl anthology wrgg//00000 202.3 emnl n. 195.
Toan Q. Nguyen and Julian Salazar. Transformers without tears: Improving the normaliza- tion of self-attention. In Proc. International Conference on Spoken Language Translation, Hong Kong, November 2019. Association for Computational Linguistics. URL https:// acl aaaahhhoooyy www.llggoyoogg/1111 9. wslt-1.17.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Py PyTorch: An imperative style, regh-perffrmance deep learning library. In Advances NewrIPS, volume 32, Vancouver, Canada, December 2019. Cur- Associates, Inc. URL https: / nips cc/ cc/papprrteeeeeeeeeeeeeeeeeeeeee files/ e /aaps//ppp222111 9/ hash/bábca228 feee wat.aeeeeeff22aaaaf1002222222200
Abst Abstract html.
Jackson Petty and Robert Frank. Transformers generalize linearly, September 2021. URL https: / hwwp:/ww.oo//ab/////1/aa 6.36. arXiv:2109.12036
12

---

Published as a conference paper at ICLR 2024
Peng Qian, Tahira Naseem, Roger Levy, and Ramón Fernandez Astudillo. Structural guidance for transformer language models. In Proc. ACL-IJCNLP (Long Papers), pp. 3735-3745, Online, August 2021. Association for Computational Linguistics. doi: 111118866110000 1111665511/2221.ccccomm-m33333 URI. https:/ / aac anntoo ogy. org/ 2021 1ar1.accl-- acl-long.22. 9.89.
Walter L. Ruzzo. On uniform circuit complexity. J. Computer and System Sciences, 22(3):365-383, URL https: /wwp://wwwwwoo. http://w.ceeeeereeeeeeec..ceeeee com/
1981. doi: 10.101600002000000811//000000000 sctienee/artccee//////2220000000
Laurent Sartran, Samuel Barrett, Adhiguma Kuncoro, Miloš Stanojević, Phil Blunsom, and Chris Dyer. Transformer grammars: Augmenting transformer language models with syntactic in- ductive biases at scale. Trans. ACI., 10:1423-1439, December 2022. doi: 10. 10.1162/tac_a. 00526. URL
mit http://c///aattie///ee/ee//////e http://ddo//11.1/..11... /tael_ttcl_aaa Transformer.
Grammars Augmenting
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proc. ACI. (Long Papers), pp. 1715-1725, Berlin, Germany, August 2016. Association for Computational Linguisties. doi: 10. 10.18663/11/11--11662 URL https: haaclannhhoooy.orrg/PP... or...11-16.1.2..1.11111111111111
Vighnesh Shiv and Chris Quirk. Novel positional encodings enable tree-based trans- formers. In Advances in NeurlPS, volume 32, Vancouver, Canada, 2019. Curran Asso- ciates, Inc. URL https: nwws.cccsppseeeee wwtpppppeepiiaaeeeeeeee 1ec91449211155100e1117ccdddee6
Michael Sipser. Introduction to the Theory of Computation. Cengage Learning, 3rd edition, 2013.
Marten Schijndel, Aaron Mueller, and Tal Linzen. Quantity doesn't buy quality syntax with neural language models. In Proc. EMNLP-IJCNLPP, pp- 5831-5837, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/11/D19-111 wrg.orrgggoo 1192299.
URL
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Ad- vances in NIPS, volume 30, Long Beach, California, December 2017. Curran Asso- ciates, Inc. URL attps::/papett www.rrsnniiii..o .ip.. cctpsppec/pppprr 3f5ee24354 11335477ee99 fbd0033 c4. oar.a22557770ee100000ooocoaaaaaa Abst ha0.0scccccca.aaaa-0aatrr....hh httles/aaper/017//aae
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F. Wong, and Lidia Chao. Learning deep transformer models for machine translation. In Proc. ACL, pp. 1810-1822, Flo- rence, Italy, July 2019a. Association for Computational Linguistics. doi: 10.18653/v1/P19-111 http:: http:///aaaaathhh ottp://atthhoooo.ooo//1/11111111
Yaushian Wang. Hung-Yt Yi Lee, and Yun-Nung Chen. Tree transformer: Integrating tree structures into self-attention. In Proc. EMNLP-IJCNLPP, pp. 1061-1070. Hong Kong, China, November 2019b. Association for Computational Linguistics. doi: 10.18653/v1/D19- 1098. URL https: //a /aaclannhhoogyy.ooooooo hwwp:/loooo.../11111//111111 1098.
Alex Warstadt, Aaron Mueller, Leshem Choshen, Ethan Wilcox, Chengxu Zhuang, Juan Ciro, Rafael Mosquera, Bhargavi Paranjabe, Adina Williams, Tal Linzen, and Ryan Cotterell. Find- ings of the BabyL.M challenge: Sample-efficient pretraining developmentally plausible cor- pora. In Proc. BabyLM Challenge at CoNLL, pp. 1-34, Singapore, December 2023. Associ- ation for Computational Linguistics. doi: 10. 1/20://1111222nllbbbbbooo homll-babylmm
URL https:
/ //aclantholoyy.... wttp:yy..rg//2330/
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention net- works can process bounded hicrarchical languages. In Proc. ACL-IJCNLP (Long Papers), pp. 3770-3785, Online, August 2021. Association for Computational Linguistics. doi: 10. 18.18v53/v// 2021 lonl 292. URL https: ac www.cllnnthooogyyo org 2021 atl.1..acloong 292.
Yian Zhang, Alex Warstadt, Xiaocheng Li, and Samuel R. Bowman. When do you need billions of words of pretraining data? In Proc. ACL-IJCNLP (Long Papers). pp. 1112-1125. Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/202/ acl-long.99. .acl 10ng.
URL https::tp/........... waw.laanhhoooyyyooooooo a00.0......11119900....
13

---

Published as a conference paper at ICLR 2024
Zhuosheng Zhang, Yuwei Wu, Junru Zhou, Sufeng Duan, Hai Zhao, and Rui Wang. SG-Net: Syntax-guide machine reading comprehension. Proc. AAAI, New York, New York, USA, February 2020. URL https:: //oss aaai. org. www.ddew.ppp//A/aatiiiiio. 6511/6367
A DETAILS OF PUSHDOWN AUTOMATA
Here, we define pushdown automata in more detail than in Section 3.2.2.
Definition 1. A pushdown automaton (PDA) tuple EQ, go, F), where
Q is finite set of states,
Σ is finite input alphabet,
a finite stack alphabet,
C Q x (EU (x)) x D* x x x
is the transition function,
4o is the start state, and
F Q is the set of accept states.
Let w = wh ... Wn E Ed be a string of input symbols. At any given time, having read up to input position PDA is in state encapsulate this as configuration (i, q, 8). The initial configuration is always (0, go. L), where LE l is a designated bottom symbol. If (g, u, v) 6, we say the PDA has transition g, 3 r, which signifies that the PDA is in state q and has u on top of the stack, and a can be scanned from the input, then it goes to state pops and pushes v.
and has stack m" (denoted in bottom-to-top order). We
run a sequence of transitions 8, linked by configurations, starting with the initial configura- tion. Since the PDA is nondeterministic, multiple transitions may ambiguously apply the same configuration, generating multiple runs per input string. We write π we q. to indicate that run i ends in configuration (i, 9, (x) for some 3 T , r. We say the PDA accepts string if there a run π ~ n, J, for some f F, and rejects it otherwise.
B IMPLEMENTATION DETAILS OF THE DVPDA
Here, we provide details of implementing the dVPDA, which consists of computing ro...., Fn (Eq. (12)) given A..... An and vo.... Un where n is the length of the input sequence. For an explanation of how this implementation is derived, see DuSell & Chiang (2020; 2023).
We maintain three main data structures:
A tensor of size (n + 1) x (n + 1) x [ x [ X Q1 x [T] called the inner weights. We write r, where 1 <
elements as y|i + q
1,0 S tt< n, q,r E Q. and x, y A tensor Ç of size (n+ + 1) x (n+ 1) x [Q] x [1' X IQ] X [1' m called the vector inner weights.
We write elements, which
members of R Ç $11g, x + (r similarly y
A tensor a size (n+ 2) x [Q] [ called the forward weights. We write elements as a |/t][r, y), where -1 < t < E Q. and r.
Let [[Ø] denote the indicator function, which is 1 if proposition & true and otherwise. We initialize the tensors as follows.
Y|-1 0 [g. r.y] I[g 9o A 1 Ar A
A =
(21)
< 00 = =I = 9o A A Ar=
(22)
a|-1][r.y] I[r 90 A
(23)
14

---

Published as a conference paper at ICLR 2024
For < < and -isist-1
y/2 + fllg, x +,
1/ 1] A,g, r,
push
Σ
H/6 + - - + 8 Ar|s, + r.
repl.
=EQ
(24)
1-2
Σ
Σ
7/e ki]1,, u 7lk + f1(u, y + pop
k= uEQ
Here, 7 is auxiliary tensor of size (t - Q x
that is computed once for each t. r [Q]
Elements are written as 7'k + (/[u,11 + r/. where
x
1. ur Q, and E r. For
k 2-
7/k - 111,,
Σ
>/k + 1- 1-111|1,,y
-8 8 A. rr,e].
(25)
sEQ,:ET
For 515 n and -1sist-1,
Ç/i + f|gg, x + r,y)
1/i = 1] Ar zz) 12
push
Σ 5/i + 1||4, + s, Arls, +r,y
repl.
=22E
(26)
2-2
Σ Σ +
f1/u, + + pop
aaitluE
For 5t5 n,
1-1
a|t[[r = Σ Σ
alill,, x] 7/ + +tlgq. + r.yl.
(27)
i=-1rEQ.EET
For each we compute tensor 7h of size [Q] [T] m. We write elements, which are members of Rmm, as [r, v), where r E Q and T.
mr y)
Σ
Σ gEQ
atillg cii fl[g, - r.vl
(28)
Finally, we compute the stack reading r., implementing Eq. (12), as follows.
mr [r, yl
r/lr vl
(29)
EreQver a[t]
To prevent underflow and overflow, we compute C, 2,, and entirely in log space.
How do these equations fit with Eq. (4)? Note that 7, Ç, and a can be computed incrementally in order of increasing t. To achieve the right time and space complexity, it is important to pre-allocate 7. Ç. and a and update them in-place. We can represent the stack state S, as (1,7, a.).. Then STACK((I 1, ax), A., vx) computes all values up to timestep assuming that all values up to 1 have been computed, and increments t to t. We let READING((t, 7, a.) implement Eqs. (28) and (29).
C TIME AND SPACE COMPLEXITY OF STACK ATTENTION
In Table 4, we show the serial time and space complexity of the three types of attention studied in this paper: SDPA, superposition stack attention (Sup), and nondeterministic stack attention (Nd). In the case of superposition stack attention, only the forward pass is needed, past stacks can be discarded at each timestep, reducing the space complexity factor of O(n).
15

---

Published as a conference paper at ICLR 2024
Table 4: Serial time and space complexity of the types of attention studied in this paper. "Serial Time" and "Space" are for a combined forward-backe backward pass; "Space (Forward Only)" is for forward pass only. Here, is the length of the imput sequence, de is the size of the query/key vectors in SDPA, de is the size of the value vectors in SDPA, m is the size of the pushed stack vectors va in stack attention, Q is the set of PDA states in the nondeterministic stack, and I is the set of stack symbols in the nondeterministic stack. For simplicity, we disregard the linear transformations of Eqs. (I) and (17) to (20).
Attention
Serial Time
Space
Space (Forward Only)
SDPA
Odgnn + don²) O(mn²)
On + +dn don)
same
Sup
(mm)²)
(mm)
Nd
Q(mQQ) T +111
(m
same
In order provide better picture of the practical speed of stack attention on GPUs, in Table 1, we show the parallel time complexity of our current implementations of stack attention, as well as improved parallel time complexity that could be attained after applying known speedups.
Stack attention introduces a recurrence that does not exist in SDPA (see Eq. (4)), so it cannot be par- allelized across the timestep dimension as easily. (Note that SDPA requires softmax over n inputs, which can be done in Ollog n) parallel time.) However, stack attention in transformers does present opportunities for parallelization that do not exist in stack-augmented RNNs. In stack-augmented RNNs, the actions and pushed vector for timestep č depend on the hidden state at timestep t. Since the hidden state computation cannot be parallelized across t, neither can the stack computation. On the other hand, in a transformer with stack attention, the stack actions and pushed vectors for all t are computed in parallel by earlier SDPA layers and are available all once. This means that we can use techniques that parallelize the stack computation across t.
In the case of nondeterministic stack attention, our current implementation, which is based on Lang's algorithm (see Appendix B), populates y and Ç serially for increasing t, incurring a cost of (()) parallel time. The update at each step is parallelized over i, but the summation over & runs (k) parallel time. So, the overall parallel time complexity is O(n?).
We can parallelize Lang's algorithm over t a manner very similar to that of parallel CKY parsing (Kosaraju, 1975). This would require iterations that populate " and Ç in increasing order of span size f-i. Each iteration would, in parallel over each i, compute the terms for all splits in the pop rule in parallel over k and them in Ollog k) parallel time. So, the overall parallel time complexity can be lowered (n log Since CFL recognition is in 2C² (Ruzzo, 1981), theoretically this can be lowered even further to ((logn))))
In the case of superposition stack attention, each stack element depends only on constant number of elements (three) from the previous timestep, so each stack update can be done in O(1) parallel time. Our implementation performs stack updates in order of increasing ((n) parallel time. Since superposition stack attention is a special case of nondeterministic stack attention, it can also be theoretically lowered to O(logn)))) ))) and perhaps further.
so overall it runs in
GPU memory limitations are sometimes a practical challenge when using nondeterministic stack attention. For the machine translation experiments in particular, which use separate stack attention modules in both the encoder and decoder. we had to contend with memory fragmentation issues in PyTorch's CUDA memory allocator.
For examples of wall-clock runtimes and GPU memory cost, see Table 6.
D DETAILS OF CONTEXT-FREE LANGUAGE EXPERIMENTS
Here, we provide additional details about the models and training procedure used in Section 6.
16

---

Published as a conference paper at ICLR 2024
Table 5: Parameter counts for each task and each model used in Section 6. across tasks due to differences in vocabulary size.
Parameter counts differ
Task
LSTM
LSTM+Sup
LSTM+Nd
Tf
Tf+Sup
Tf+Nd
42,004
41,402
31,138
43,044
40,964
33,273
www
41,503
40,936
30,817
42,979
40,899
33,216
woPw
41,503
40,936
41,482
42,979
40,899
36,576
Dyck
42,505
41,868
31,459
43,109
41,029
33,330
Hardest CFL
44,008
43,266
43,087
43,304
41,224
36,861
D.1 MODELS
For LSTMs, inputs are encoded as one-hot vectors. We PyTorch's (Paszke al., 2019) LSTM implementation, which by default includes redundant bias parameters bbi dht, big. and bho- We remove these parameters in our experiments and from our reported parameter counts. We set the initial hidden state and memory cell to O.
For transformers, we use PyTorch's transformer layer implementation. Like Vaswani et al. (2017), we map inputs to vectors of size dmodel with a scaled embedding layer and apply sinusoidal posi- tional encodings. We do not tie input and output embeddings. The first input token to each trans- former is always BOS. We train all models to predict EOS as the last output token. We map the outputs of the final layer to logits for predicting the next token via learned affine transformation. We use dropout rate of 0.1 throughout the transformer; PyTorch applies it in the same places as Vaswani et al. (2017), and additionally to the hidden units feedforward sublayers and the atten- tion probabilities of SDPA (z, in Eq. (2)). We apply dropout to the hidden units of the feedforward sublayers of stack attention layers as well.
We show parameter counts for all models on all tasks in Table 5.
D.2 INITIALIZATION
We initialize all fully-connected layers with Xavier uniform initialization (Glorot & Bengio, 2010) except for the recurrent L.TTM layer and all fully-connected layers in SDPA layers. For layer norm, we initialize all weights to I and all biases to O. We initialize all other parameters by sampling uniformly from [-0.1,
D.3 TRAINING
We use minibatches of size 10. We generate batches once before training: to generate batch, we first select length and then sample strings of that length. We randomly shuffle batches before each epoch. We train each model by minimizing its cross-entropy (summed over the timestep dimension of each batch) on the training set, and we use per-symbol cross-entropy on the validation set as the early stopping criterion. We optimize parameters with Adam (Kingma & Ba, 2015), and randomly sample the initial learning rate from a loo-lloiilte uniform distribution over [5 x 10 10-4, x 10 2). We clip gradients with threshold of 5 using 2 norm rescaling. We multiply the learning rate by 0.9 after 5 epochs of no improvement in cross-entropy on the validation set, and we stop early after 10 epochs of improvement. We train for maximum of 200 epochs.
E ANALYSIS OF STACK ATTENTION
In Figs. 3 and 4, we show heatmaps of the stack actions learned by T the real-time DCFL. tasks. These actions are from the same models whose results are shown Fig.
F DETAILS OF NATURAL LANGUAGE MODELING EXPERIMENTS
Here, provide additional details about the models and training procedure used in Section 7.
17

---

Published as a conference paper at ICLR 2024
0
1
0
1
1
0
1
0
0
0
1
1
1
1
1
0.8
0.8
0
f Symbol
0.6
Probability
< Symbol
0.6
Probability
0
0
0.4
0.4
0
0.2
0.2
1
0
0
1
1
0
0
1
0
1
1
0
1
0
EOS
EOS
push
no-op pop
push no-op pop
Action
Action
Figure 3: Visualization of superposition stack attention on a string in wt " R As ex- pected, the model learns to push all symbols before 4, do nothing when reading #, and pop all symbols after 4.
Figure 4: Visualization of superposition stack attention a string in the Dyck lan- guage. As expected, the model learns to push opening brackets and pop when reading clos- ing brackets.
18

---

Published as a conference paper at ICLR 2024
Table 6: Computational cost of training each architecture the PTB language modeling task when run on an NVIDIA TITAN Xp GPU.
Model
Examples/s
Minutes/Epoch
GPU Memory
Tf
859
0.8
394 MB
Tf+Sup
345
1.9
397 MB
Tf+Nd
27
24.3
1.91 GB
F.1 MODELS
The first input token each transformer is always BOS, and train predict EOS the last output token. We use dropout rate of 0.1 in all the same places as in Appendix D.1. We use sinusoidal positional encodings. We do tie the embeddings of the input and output layers.
F.2 INITIALIZATION AND TRAINING
We initialize parameters in the same way as Appendix D.2, except the interval used for uniform sampling is 0-0.1 0.01].
For each run, we randomly sample a batch size B from a uniform distribution over [128, 512]. For each epoch, we randomly shuffle examples and group examples of similar lengths into the same minibatch, enforcing an upper limit of tokens per batch, including padding, BOS, and EOS tokens. We clip gradients with a threshold of using L norm rescaling.
optimize parameters using Adam. For each run, we randomly sample initial learning rate from log-uniform distribution over [10-6, 10-4). We take checkpoint every 20k examples to evaluate the model's 's cnoss-entropy on the validation set. We multiply the learning rate by 0.5 after checkpoints with no improvement on the validation set, and stop early after 4 checkpoints with no improvement. We use the checkpoint with the lowest validation cross-entropy.
We show the computational cost of training each architecture this task in Table 6.
G DETAILS OF MACHINE TRANSLATION EXPERIMENTS
Here, we provide additional details about the data preprocessing, models, training procedure, and decoding algorithm used in Section 8.
G.1 DATA PREPROCESSING
We limit the training. validation, and test examples where the source and target side each have no more than 150 Unicode characters after applying NFKC normalization. In the training set, we filter training examples with empty sentences and those where the length in characters of one side is more than 4 times that of the other. After filtering, we train a tokenizer on the combined source and target sentences of the subsampled training data using the SentencePiece (Kudo & Richardson, 2018) implementation of BPE (Sennrich et al., 2016). We use a vocabulary of 32k tokens and frequency threshold of 50.
G.2 MODELS
We always give the encoder BOS as the last input token to allow to detect the end of the source sequence. We always give the decoder BOS as the first input token and have predict EOS as the last output token. SDPA is not causally masked in the encoder, but it is in the decoder. Like Vaswani et al. (2017), tie the embeddings of the encoder input, decoder input, and decoder output layers. We use a dropout rate of 0.1 as in Appendix F.I, applying it also to the softmax probabilities of cross-attention.
19

---

Published as a conference paper at ICLR 2024
G.3 INITIALIZATION AND TRAINING
We initialize parameters in the same way as Appendix F.2. For every epoch, we randomly shuffle the training data and group examples where the combined source and target length in tokens similar into the same minibatch. We limit the number of tokens in the source or target side of a batch to 2048 tokens, including padding, BOS, and EOS symbols. We use label smoothing with weight of 0.1. We clip gradients with a threshold of using L norm rescaling.
We optimize parameters using Adam. We randomly sample the initial learning rate from log- uniform distribution over [10-6, 10-3). We take checkpoint every 50k examples to evaluate the decoder's cross-entropy on the validation set. We multiply the learning rate by 0.5 after 2 check- points with no improvement on the validation set, and we stop early after 4 checkpoints with no improvement. We train for maximum of 100 epochs. We use checkpoint with the best valida- tion cross-cntropy.
G.4 DECODING
We use beam search for decoding, using a beam size of 4. During beam search, we apply length nor- malization to the probabilities of hypotheses before selecting the top k hypotheses for the next beam. We do this by dividing the oog-robbayility of each hypothesis by the number of tokens generated by that hypothesis so far (including EOS).
20


Reviews:
[{"review_id": "asOvuWvqk6", "reviewer": "Reviewer_ecSr", "summary": "This paper presents a stack-augmented transformer to help overcome some of the challenges of ordinary self-attention in modeling nested syntactic structure. \n\nThe overall approach is to propose two different stack mechanisms: the \"superposition stack\" which is an extension of Joulin and Mikolov (2015), and a second non-deterministic diferentiable vector pushdown automata (dVPDA). While the details for the dVPDA are not very clear from the methods section (see Weaknesses and questions), both the dVPDA and the superposition stack provide soft-stack vector readouts at each time-step, and can be used to replace standard attention. \n\nFrom results, we find:\n- On formal languages such as Dyck, 5 layer transformers with the stack based attention (coming from the superposition stack) are better than ordinary 5 layer transformers (though both of these are worse than LSTM variants). \n- On language modeling over the penn treebank, we find that transformers + dVPDA obtain lower perplexities than ordinary transformers (though transformers + superposition stack are much worse). \n- On a 5 layer machine translation dataset, we find slight very mixed improvements in BLEU.", "questions": "I'll intersperse questions (Q) with some suggestions for improving writing (S)\n\nIntroduction:\n\nQ1: \"Recent work has shown that transformers have linear rather than hierarchical..\": Is this true? Murty 2023 (\"Grokking\") find that transformers acquire hierarchical bias when trained for long. It would be better to qualify this statement somewhat.\n\nS1: \"on a natural language modeling benchmark\"  => would be better to just say \"Penn TreeBank\". In general, I found the last para to be very vague. It would be better to have concrete numbers and dataset names.\n\nRelated Work:\n\nS2: Missing several keys papers (Ordered Memory, Unsupervised Tree LSTMs, RL-SPINN) and early works on incorporating stack mechanism into transformers (Das, 1993). This is just a quick list, but there is a long history of learning syntax unsupervisedly / augmenting neural models with stacks that is completely missing.\n\nBackground Section 3.2.1:\n\nQ3: This method seems like it would run faster than the method from 3.2.2. Could the authors confirm this?\n\nQ4: In natural language, one might want to do multiple reduce operations after emitting a word. How does this approach allow for multiple reduces at each time step? \n\nBackground Section 3.2.2:\n\nS3: Definition 1 and the next paragraph take too much space and seem like background that could be in an appendix. \n\nQ5: Missing detail: What is the time complexity of Eq. 17 -  Please include pseudocode here.\n\nQ6: \"Each $a_t$ is a flattening of tensor...\": How was the size of the tensor $\\Delta_t$ computed here?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The motivation behind this paper is great - there are clear limitations of self-attention w.r.t. modeling syntactic patterns and here, we see a novel approach to use a stack to model such patterns.", "weaknesses": "- Unfortunately, I think the results are very mixed, experiments are done on very small 5 layer transformers on small datasets.\n- Even for the positive results, there is no analysis of why the stack augmented model may be doing better on natural language - is it discovering good parses / something else?\n\n-  The exposition is very confusing, and i'm a bit lost on various details. The biggest missing detail is how training is done in parallel - from Eq 8, 9 ..., 17 and Figure-1 it seems like there is a stack state that is recurrently updated, but transformer training is fully parallel, so how can state information be passed between different tokens? Is this done by basically reconstructing the previous state of the stack since all previous actions are available to the model at each time step - if so what is the FLOP hit from doing this?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AXUa23JrlN", "reviewer": "Reviewer_knG9", "summary": "The authors present work on stack attention for transformers that attempts to address naturally hierarchically structured problems. They carefully present scaled dot-product attention, a core component of the transformer architecture. Next, they provide background on differentiable stacks, superposition stacks, and a non-deterministic stack -- the differentiable vector push-down automaton. Superposition stacks can be seen as a special case of this. This is used to replace scaled dot-product attention.\n\nThey explore the empirical performance of this approach on a range of tasks, including constructed languages, a small scale language modeling problem, and a small scale machine translation effort. In some settings, the Tf-Nd configuration outperforms a conventional transformer architecture. However, these gains are not huge, the datasets are small, and (in machine translation) as the dimensionality increases, the architecture no longer shows gains.", "questions": "What is the empirical cost of running these methods?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors carefully present their formalism, including details and relevance connections to prior work.\n\nThe description of the method stands alone reasonably well, though familiarity with related work makes the paper much more accessible.\n\nThe authors evaluate in a range of settings, from constructed language to actual use cases.\n\nThe methodology does not rely on a specific grammar or automaton structure; instead it is latent.", "weaknesses": "The evaluation settings are quite small by modern standards.\n\nGains (at least in machine translation) seem to disappear as model size increases.\n\nComputational cost at inference seems greater according to the big-O runtimes -- is this a net win? Is it feasible to train at larger scale?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "arLLsW5sCB", "reviewer": "Reviewer_bkRf", "summary": "This paper addresses the problem of the lack of modeling the hierarchical structure of input sequences in transformers. In response, a new attention mechanism, called stack attention, is proposed. The idea is to frame the sequence modeling problem as a task of running a PDA on the input vectors; that is, each time we take a vector as input, we update the stack of the PDA and produce a new vector as output. In this way, the sequence of input vectors is represented as a sequence of output vectors. The entire process is based on differentiable states and operations. Thus, this attention model operates in the same manner as those in sequence modeling, making it easy to incorporate the model into transformers. However, as a side effect, stack attention introduces recurrence into modeling, thus preventing training from being parallelized across the sequence. The stack attention models are tested on synthetic data generated by context-free grammars. Experimental results show that nondeterministic stack attention models surpass standard transformers and achieve better results than a strong nondeterministic stack RNN baseline. The stack attention models also show promising results on small-scale machine translation and language modeling tasks.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I like this work! Given that language structure is not explicitly modeled in current Transformer models, this work opens a door to a new approach to considering hierarchical patterns in modeling languages. The design of the model is simple and elegant. The experiments support the claims well.", "weaknesses": "I have no major concerns, but a few comments.\n\nThe state of the stack attention models at a given timestep depends on its past state. This means the models share similar drawbacks and merits with recurrent models like RNNs. Compared to the self-attention used in standard Transformers, stack attention is slower for training because it processes one token at a time, rather than parallelizing the encoding process over the entire sequence. The author states in the appendix that this could be improved using the parallel prefix sum method, but no details are presented.\n\nA related problem is that the experiments here are small-scale. While it's fine to test the models on synthetic data for CFL tasks, the results on language modeling and machine translation aren't comparable to those in other papers. I understand that computational cost is a concern. However, to demonstrate the superiority of stack attention, it's necessary to compare it with previously reported results under the same setup.\n\nI’m not quite satisfied that the models are motivated by handling hierarchical structure behind languages but there is no discussion on what structure is captured. A simple way to examine this is to design probing tasks to see how much syntax is modeled in stack attention and to see how the learned syntax differs from human-annotated syntax. Unsupervised learning of syntactic structures can offer new insights into modeling natural languages.\"\n\nThere have been previous studies on extending standard attention models to hierarchical models, such as hierarchical attention and selective attention. These should be considered baselines for comparison, either in related work or experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)