Paper: Stack Attention: Improving the Ability of Transformers to Model Hierarchical Patterns (ICLR-2024)

Reviewer comment (Experiments):
Unfortunately, I think the results are very mixed, experiments are done on very small 5 layer transformers on small datasets.

Author response:
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).

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Reproducibility):
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?

Author response:
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.

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

Reviewer comment (Writing):
Q1: "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.

Author response:
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.

Reviewer comment (Experiments):
S1: "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.

Author response:
S1: We will follow this suggestion. Thank you!

Reviewer comment (Novelty):
S2: 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.

Author response:
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.

Reviewer comment (Evaluation):
Q3: This method seems like it would run faster than the method from 3.2.2. Could the authors confirm this?

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

Reviewer comment (Theory):
Q4: 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?

Author response:
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).

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

Reviewer comment (Presentation):
S3: Definition 1 and the next paragraph take too much space and seem like background that could be in an appendix.

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

Reviewer comment (Reproducibility):
Q6: "Each $a_t$ is a flattening of tensor...": How was the size of the tensor $\Delta_t$ computed here?

Author response:
Q6: Thank you, we should have explained this better. Let us write the weight of transition $q, x \xrightarrow{w_t} r, v$ as $\Delta_t[q, x \rightarrow r, v]$, where $q, r \in Q$ are states, $x \in \Gamma$ is the popped stack symbol, and $v \in \Gamma^\ast$ is the string of pushed stack symbols. The first three dimensions of $\Delta_t$, $|Q| \times |\Gamma| \times |Q|$, correspond to the variables $q, x, r$. According to Equations 14-16, the string $v$ has only three forms, where $y \in \Gamma$:

1. push: $xy$
2. replace: $y$
3. pop: $\varepsilon$

So there are $|\Gamma| + |\Gamma| + 1$ possible choices for $v$, which is why the last dimension is $(2|\Gamma| + 1)$.