Stack Attention: Improving the Ability of Transformers to Model Hierarchical Patterns

Dear reviewer bkRf,

Thank you very much for your review and your many thoughtful comments. We are glad that you liked the paper! We greatly appreciate your suggestions for improving it further.

The 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.

Your comments about the recurrence introduced by the stack are exactly correct. Actually, stack attention, despite having a recurrence across the timestep dimension, can still parallelize the stack computation in a way that is impossible in an RNN under the controller-stack paradigm used by Grefenstette et al. (2015), Joulin and Mikolov (2015), DuSell and Chiang (2020), etc. In an RNN, the stack actions for timestep $t$ depend on the hidden state at timestep $t$. Since the hidden state computation cannot be parallelized across $t$, neither can the stack. However, in a transformer with stack attention, the stack actions $\mathbf{a}_t$ and pushed vectors $\mathbf{v}_t$ for all timesteps are computed in parallel and are available all at once. This means we can take advantage of algorithms that parallelize the stack computation over the timestep dimension, such as parallel CKY. (Apologies, we realized that the algorithm for parallelizing nondeterministic stack attention across the timestep dimension would resemble parallel CKY, not parallel prefix sum as stated in Appendix B. We will update Appendix B.)

For nondeterministic stack attention, the algorithm would work by updating the dynamic programming table (see Appendix A) in increasing order of span size, then parallelizing the inner loops that implement the replace and pop rules. This would reduce the runtime to $O(n)$ parallel time. As for the superposition stack, each recurrent step runs in $O(1)$ parallel time under our current implementation, so it already runs in $O(n)$ parallel time.

I’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."

Thank you for raising this point. We agree that analyzing the parses learned by stack attention would be extremely informative and would further validate our method. If we have time, we will update our response with an analysis of how the networks learn to use stack attention.

There 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.

Thank you for this suggestion. To clarify, are you referring to these papers?

• Hierarchical Attention Networks for Document Classification (Yang et al.)
• Selective Attention for Context-aware Neural Machine Translation (Maruf et al.)

Dear reviewer knG9,

Thank you very much for your thoughtful review and positive comments.

The evaluation settings are quite small by modern standards.

Gains (at least in machine translation) seem to disappear as model size increases.

It's true that the models and datasets are relatively small compared to some other papers, but we believe our experiments demonstrate an important point: that transformers with a latent model of syntax are more expressive than standard transformers (as shown by the results on synthetic CFLs such as the hardest CFL), and that they can be more parameter- and data-efficient than standard transformers when the data can be explained with underlying hierarchical rules (as shown by the language modeling results).

Computational cost at inference seems greater according to the big-O runtimes -- is this a net win? Is it feasible to train at larger scale?

What is the empirical cost of running these methods?

We discuss the computational cost of stack attention in detail in Appendix B, including wall-clock runtimes on the PTB language modeling task and asymptotic time/space complexity. With our current implementation, large-scale training is probably within reach for superposition stack attention but not yet for nondeterministic stack attention; superposition stack attention trains at 40% the speed of standard attention, whereas nondeterministic stack attention trains at about 3% speed. However, we can likely improve the runtime of nondeterministic stack attention substantially by parallelizing the dVPDA computation across the timestep dimension using a strategy similar to that of parallel CKY, which would reduce it to only $O(n)$ parallel time. We can also get a further, smaller speedup by parallelizing the computation of the stack reading $\mathbf{r}_t$ from $\mathcal{S}_t$.

Dear reviewer ecSr,

Thank you very much for your thoughtful review. We greatly appreciate your constructive comments and your suggestions for improving the clarity of the paper.

Unfortunately, I think the results are very mixed, experiments are done on very small 5 layer transformers on small datasets.

It's true that the models and datasets are relatively small compared to some other papers, but we believe our experiments demonstrate an important point: that transformers with a latent model of syntax are more expressive than standard transformers (as shown by the results on synthetic CFLs such as the hardest CFL), and that they can be more parameter- and data-efficient than standard transformers when the data can be explained with underlying hierarchical rules (as shown by the language modeling results).

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?

Thank you for raising this point. We agree that analyzing the parses learned by stack attention would be extremely informative and would further validate our method. If we have time, we will update our response with an analysis of how the networks learn to use stack attention.

The biggest missing detail is how training is done in parallel ...

You are correct that the stack state $\mathcal{S}_t$ is updated recurrently. In our implementation, we simply compute $\mathcal{S}_1, \mathcal{S}_2, \ldots$ serially and compute the other parts of the transformer network in parallel as usual. Consider Figure 1, which depicts a single stack attention sublayer: we first compute all of the inputs $\mathbf{x}_t$ and apply layer norm in parallel, then we compute the stack states $\mathcal{S}_t$ and their stack readings $\mathbf{r}_t$ serially, and finally we apply dropout and residual connections in parallel to compute the outputs $\mathbf{y}_t$. Although our implementation does not currently do so, it would also be possible to parallelize the computation of $\mathbf{r}_t$ from $\mathcal{S}_t$, which would result in a small speedup for nondeterministic stack attention.

Despite the recurrence of $\mathcal{S}_t$, in the case of nondeterministic stack attention, it is still possible to parallelize the computation across the timestep dimension using a strategy similar to that of parallel CKY parsing, which would reduce it to $O(n)$ parallel time and likely result in a substantial speedup over our current implementation.

Q1: Thank you for pointing this out; this is indeed a very relevant finding. We will update this statement in our paper. It is interesting to note that although training for longer does improve hierarchical generalization by a lot, vanilla transformers still fall 8 to 20 percent short of perfect on question formation and tense reinflection, indicating there is still room for improvement.

S1: We will follow this suggestion. Thank you!

S2: Thank you for these references. Incorporating stacks and latent syntax into neural networks has become a very long line of work, so we limited our citations specifically to syntax-oriented transformers (as opposed to RNNs), and RNNs that use the differentiable stacks that we use in this paper. If we have space, we will include more discussion of other work.

Q3: Yes, superposition stack attention is faster than nondeterministic stack attention. For details, please see Appendix B.

Q4: Superposition stack attention does not appear to have a mechanism for doing this. As discussed in Section 4, superposition stack attention is very likely less expressive than nondeterministic stack attention, which is at least weakly equivalent to any PDA that performs multiple reductions per timestep, since it can recognize all CFLs (DuSell and Chiang, 2023).

S3: Thank you; we will consider ways to make this more concise while still establishing the essential notation used in the paper.

Q5: The purpose of Equation 17 is to define the output of the dVPDA mathematically in a way that is concise, if inefficient. If implemented naively, it would take time exponential in $n$, because $\pi$ loops over sequences of length $n$ in the denominator. However, it is possible to implement this equation in cubic time and quadratic space using a dynamic programming algorithm called Lang's algorithm (this is the same time/space complexity as parsing algorithms for context-free grammars, such as CKY). Since Equation 17 is sufficient for understanding the dVPDA mathematically and the implementation details are somewhat complex, we put them in the appendix to save space and to reduce the cognitive burden on the reader. For the full implementation details of Equation 17, please see Appendix A, and for a discussion of its time complexity, please see Appendix B.

(continued below)

Re: parallel vs recurrent, based on your explanation, it appears that only layers below the first stack attention layer are parallelizable while the rest of the model is still recurrent?

The layers after stack attention remain parallelizable as well. The procedure is essentially:

1. compute all layers before stack attention in parallel to get $\mathbf{a}_t$ and $\mathbf{v}_t$ for all $t$
2. run stack attention serially to get $\mathbf{r}_t$ for all $t$
3. now that we have all $\mathbf{r}_t$, compute all $\mathbf{y}_t$ (see Figure 1) and all layers after stack attention in parallel

I would also still like to know the FLOP hit from this recurrent computation.

We have not yet worked out the implementation details of the parallelized version, so at this time it's difficult to quantify the difference in FLOPs. In terms of time complexity, the current implementation runs in $O(n^2)$ parallel time, whereas the parallel CKY-style version would run in $O(n)$ parallel time.

And additionally, could you also please confirm the addition parameters that are added in stack attention?

The SDPA sublayer function uses the parameters $\mathbf{W}\_{\mathrm{q}}, \mathbf{W}\_{\mathrm{k}} \in \mathbb{R}^{d\_k \times d\_{\mathrm{model}}}$, $\mathbf{W}\_{\mathrm{v}} \in \mathbb{R}^{d_v \times d\_{\mathrm{model}}}$, and $\mathbf{W}\_{\mathrm{y}} \in \mathbb{R}^{d\_{\mathrm{model}} \times d\_v}$. Stack attention replaces these with the parameters $\mathbf{W}\_{\mathrm{a}} \in \mathbb{R}^{d\_a \times d\_{\mathrm{model}}}$, $\mathbf{W}\_{\mathrm{v}} \in \mathbb{R}^{m \times d\_{\mathrm{model}}}$, and $\mathbf{W}\_{\mathrm{y}} \in \mathbb{R}^{d\_{\mathrm{model}} \times d\_r}$. Nondeterministic stack attention additionally adds $\mathbf{w}\_{\mathrm{v}} \in \mathbb{R}^m$, which is used for learning the initial bottom vector $\mathbf{v}\_0$. In the case of superposition stack attention, $d\_a = 3$ and $d\_r = m$. In the case of nondeterministic stack attention, $d\_a = |Q| \cdot |\Gamma| \cdot |Q| \cdot (2|\Gamma| + 1)$ and $d\_r = |Q| \cdot |\Gamma| \cdot m$.

The superposition stack works much better than the non-deterministic stack (both in terms of speed and results on natural language data) but as you said, it is unable to make multiple reduces at each token. Does this then mean that it cannot produce natural language syntax trees?

You are correct that because of superposition stack attention's inability to perform multiple reductions per timestep, we expect that it would not be effective at modeling certain syntactic patterns that require closing multiple constituents at once, such as prepositional phrase attachment. However, it should still be advantageous for a subset of syntactic patterns where this is not an issue, such as center embedding (e.g. "The diamond that the thief who the man saw stole glittered") and gross syntactic expectation (e.g. an independent clause must come after a subordinate clause beginning with "As"). Note that although superposition stack attention does perform better than nondeterministic stack attention on our MT task, nondeterministic stack attention has better perplexity on language modeling. We think this is because of the difference in expressivity.