# Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study

arXiv: https://arxiv.org/abs/2410.17980 · HTML: https://mlsanity.com/paper/2410.17980/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 51 references that are in the map.
- First posted on arXiv 23 Oct 2024.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [RACE: Large-scale ReAding Comprehension Dataset From Examinations](https://mlsanity.com/paper/1704.04683/) (2017)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [TinyLlama: An Open-Source Small Language Model](https://mlsanity.com/paper/2401.02385/) (2024)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Zoology: Measuring and Improving Recall in Efficient Language Models](https://mlsanity.com/paper/2312.04927/) (2023)
- [The Illusion of State in State-Space Models](https://mlsanity.com/paper/2404.08819/) (2024)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing](https://mlsanity.com/paper/2306.12929/) (2023)
- [Transformer Language Models without Positional Encodings Still Learn Positional Information](https://mlsanity.com/paper/2203.16634/) (2022)
- [Language Modeling with Deep Transformers](https://mlsanity.com/paper/1905.04226/) (2019)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)
- [The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization](https://mlsanity.com/paper/2110.07732/) (2021)
- [Neural Language Modeling by Jointly Learning Syntax and Lexicon](https://mlsanity.com/paper/1711.02013/) (2017)

## Built on by (most cited first, 13 of 13)

- [Attention Residuals](https://mlsanity.com/paper/2603.15031/) (2026)
- [Forgetting Transformer: Softmax Attention with a Forget Gate](https://mlsanity.com/paper/2503.02130/) (2025)
- [PaTH Attention: Position Encoding via Accumulating Householder Transformations](https://mlsanity.com/paper/2505.16381/) (2025)
- [Long-Context Generalization with Sparse Attention](https://mlsanity.com/paper/2506.16640/) (2025)
- [Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings](https://mlsanity.com/paper/2509.10534/) (2025)
- [Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access](https://mlsanity.com/paper/2504.16795/) (2025)
- [On Vanishing Variance in Transformer Length Generalization](https://mlsanity.com/paper/2504.02827/) (2025)
- [Understanding and Improving Length Generalization in Hierarchical Sparse Attention Models](https://mlsanity.com/paper/2510.17196/) (2025)
- [Agentic Discovery of Neural Architectures: AIRA-Compose and AIRA-Design](https://mlsanity.com/paper/2605.15871/) (2026)
- [From Next Token Prediction to (STRIPS) World Models](https://mlsanity.com/paper/2509.13389/) (2025)
- [An empirical study on the limitation of Transformers in program trace generation](https://mlsanity.com/paper/2509.25073/) (2025)
- [Beyond Static and Linear: What Attention Constraints Best Fit Human Reading Times?](https://mlsanity.com/paper/2608.23818/) (2026)
- [How Local Mixing Encodes Relative Position in Global NoPE Attention](https://mlsanity.com/paper/2609.38109/) (2026)

## Abstract (arXiv)

The self-attention mechanism traditionally relies on the softmax operator, necessitating positional embeddings like RoPE, or position biases to account for token order. But current methods using still face length generalisation challenges. We investigate an alternative attention mechanism based on the stick-breaking process in larger scale settings. The method works as follows: For each token before the current, we determine a break point, which represents the proportion of the stick, the weight of the attention, to allocate to the current token. We repeat this on the remaining stick, until all tokens are allocated a weight, resulting in a sequence of attention weights. This process naturally incorporates recency bias, which has linguistic motivations for grammar parsing. We study the implications of replacing the conventional softmax-based attention mechanism with stick-breaking attention. We then discuss implementation of numerically stable stick-breaking attention and adapt Flash Attention to accommodate this mechanism. When used as a drop-in replacement for current softmax+RoPE attention systems, we find that stick-breaking attention performs competitively with current methods on length generalisation and downstream tasks. Stick-breaking also performs well at length generalisation, allowing a model trained with 211 context window to perform well at 214 with perplexity improvements.