Research map / Architectures and efficiency / Attention and transformers
Stack Attention: Improving the Ability of Transformers to Model Hierarchical Patterns
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes stack attention, replacing SDPA sublayers with differentiable stacks (superposition and nondeterministic vector PDA) so transformers can model hierarchical/CFL patterns without syntactic supervision.
Problem
Standard scaled dot-product attention lacks a mechanism for hierarchical patterns of arbitrary nesting depth, limiting transformers on context-free syntactic structures and making them data-inefficient for language.
Approach
Replace the SDPA sublayer with a differentiable stack: soft-selection over sequences of push/pop/no-op actions (superposition stack) or over VPDA runs (nondeterministic dVPDA), returning the top stack vector, enabling a causal latent syntax model trained end-to-end.
Key results
- On Penn Treebank under constrained parameters, Tf+Nd stack attention reaches 88.54 test perplexity versus 92.84 for the SDPA transformer, with the fewest parameters (9.86M vs 10.05M).
- Nondeterministic stack attention achieves the best in-distribution Hardest CFL performance, outperforming all models including the VRNS-RNN, despite having the fewest parameters.
- On machine translation (100k Europarl subset), stack attention does not consistently improve over the SDPA transformer; best BLEU is 12.69 for baseline Tf at the largest tier.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace one attention sublayer with superposition stack attention | Tf · Penn Treebank | perplexity: 115.11 -> 122.94 | Hurt |
| replace one attention sublayer with nondeterministic stack attention | Tf · Penn Treebank | perplexity: 115.11 -> 110.59 | Helped |
| replace one attention sublayer with superposition stack attention | Tf · Penn Treebank test | perplexity: 92.84 -> 98.67 | Hurt |
| replace one attention sublayer with nondeterministic stack attention | Tf · Penn Treebank test | perplexity: 92.84 -> 88.54 | Helped |
| use nondeterministic stack attention instead of standard attention | Tf · UnmarkedReversal, PaddedReversal, and Hardest CFL | cross-entropy difference | Helped |
| use nondeterministic stack attention instead of standard attention | Tf · Hardest CFL, in-distribution lengths | cross-entropy difference | Helped |
| use nondeterministic stack attention instead of standard attention | Tf · MarkedReversal and Dyck, out-of-distribution lengths | cross-entropy difference | Hurt |
| use superposition stack attention instead of standard attention | Tf · Penn Treebank | perplexity: 115.11 -> 122.94 | Hurt |
| use superposition stack attention instead of standard attention | Tf · Europarl de-en newstest2016, d_model 240 | perplexity: 12.54 -> 11.53 | Helped |
| use superposition stack attention instead of standard attention | Tf · Europarl de-en newstest2017, d_model 240 | BLEU: 12.11 -> 12.81 | Helped |
| use superposition stack attention instead of standard attention | Tf · Europarl de-en newstest2017, d_model 360 | BLEU: 12.69 -> 12.03 | Hurt |
| use nondeterministic stack attention instead of standard attention | Tf · Europarl de-en newstest2017, d_model 240 | BLEU: 12.11 -> 11.50 | Hurt |
Surprising findings
- Superposition stack attention does not outperform the baseline transformer on Penn Treebank despite being a special case of the dVPDA.
- Nondeterministic stack attention is about 30x slower in training throughput than SDPA (27 vs 859 examples/s).
Limitations
- Nondeterministic stack attention is computationally costly (O(n^2) or cubic time implementation) and memory-hungry.
- Machine translation shows no consistent gains over baseline transformers.
- Transformers with stack attention still show length generalization issues on DCFLs.
Experimental setup
Datasets: Marked Reversal Unmarked Reversal Padded Reversal Dyck Hardest CFL Penn Treebank Europarl v7 de-en newstest2016 newstest2017
Benchmarks: Marked Reversal Unmarked Reversal Padded Reversal Dyck Hardest CFL Penn Treebank newstest2017
Metrics: cross-entropy difference (nats) perplexity BLEU
Baselines: Transformer with SDPA Tf+Sup (superposition stack attention) LSTM LSTM+Sup VRNS-RNN (LSTM+Nd) Transformer-XL
Builds on: superposition stack dVPDA (Dusell and Chiang 2023) stack-augmented LSTM scaled dot-product attention
Closest papers by meaning
- Core Context Aware Transformers for Long Context Language Modeling (ICML 2025 · similarity 0.87)
- Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study (ICLR 2025 · similarity 0.86)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.86)
- Punctuation-aware Hybrid Trainable Sparse Attention for Large Language Models (NeurIPS 2026 · similarity 0.85)
- Training-Free Long-Context Scaling of Large Language Models (ICML 2024 · similarity 0.85)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.