Research map / Architectures and efficiency / Attention and transformers
Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study
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 stick-breaking attention as a drop-in softmax replacement without positional embeddings, with a numerically stable Triton/FlashAttention-style kernel, and evaluates length generalisation and 1B/3B pretrained models.
Problem
Softmax attention needs positional embeddings (e.g., RoPE) and existing position methods generalise poorly to longer contexts than seen in training.
Approach
Replace softmax with a stick-breaking process where each token's weight is a sigmoid break point applied to the remaining stick, giving an inherent recency bias and removing the need for position embeddings; implement it with a log-space, block-wise Triton kernel adapted from FlashAttention-2.
Key results
- On 1B models, stick-breaking averages 59.5 vs 58.8 for softmax on the NLP suite and gets Wikitext perplexity 13.4 vs 13.8.
- On MMLU 0-shot, the 3B stick-breaking model scores 50.8 vs 46.1 for softmax; on 8-shot CoT GSM8K it scores 49.7 vs 44.2.
- On 350M models trained at 1024 context, stick-breaking NLL keeps decreasing as context grows to 64K while other position methods plateau or rise.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Softmax + RoPE attention with stick-breaking attention | 1B pretrained model · NLP benchmarks | Avg. accuracy: 58.8 -> 59.5 | Helped |
| Replace Softmax + RoPE attention with stick-breaking attention | 1B pretrained model · WikiText | perplexity: 13.8 -> 13.4 | Helped |
| Replace Softmax + RoPE attention with stick-breaking attention | 3B pretrained model · NLP benchmarks | Avg. accuracy: 63.3 -> 64.1 | Helped |
| Replace Softmax + RoPE attention with stick-breaking attention | 3B pretrained model · WikiText | perplexity: 11.3 -> 10.8 | Helped |
| Replace Softmax + RoPE attention with stick-breaking attention | 3B model · GSM8K 5-shot | accuracy: 44.1 -> 42.3 | Hurt |
| Replace Softmax + RoPE attention with stick-breaking attention | 3B model · GSM8K 8-shot, CoT | accuracy: 44.2 -> 49.7 | Helped |
| Use stick-breaking attention instead of Softmax + RoPE | 2-layer Transformer · MQRAR | maximum key-value pairs handled: 128 -> 192 | Helped |
| Use the Triton stick-breaking attention kernel instead of FA2 | 1B class training · Training throughput | billion tokens / day: 29.5 billion tokens / day -> 21 billion tokens / day | Hurt |
| Enable forward-pass early halting / block skipping (ablation) | Stick-breaking attention · Pile 10k, 16K context | speed improvement: 9.3% | Helped |
| Add remainder bias (ablation) | Stick-breaking Transformer · 1.2B parameters, 1.25T tokens | Avg.: 47.2 -> 47.5 | Helped |
| Add Group Norm to stick-breaking with remainder bias (ablation) | Stick-breaking Transformer · 1.2B parameters, 1.25T tokens | Avg.: 47.5 -> 47.5 | No clear effect |
| Add remainder bias and Group Norm (ablation) | Stick-breaking Transformer · 1B activation / 6.5B parameters, 2T tokens | Avg. | Helped |
Surprising findings
- NoPE loss increased as context length was extended, contrary to prior reports that NoPE generalises well.
- Stick-breaking Transformers improved on softmax with the remainder bias and group norm additions even though these were added for stability reasons.
Limitations
- Triton kernel is about 29% slower than FlashAttention-2 in throughput.
- Attention weights do not sum to one, requiring a remainder bias workaround for full benefit.
- Length extrapolation experiments are on small 350M models and RULER evaluation is on models trained only at 4096 context.
Experimental setup
Datasets: SlimPajama The Pile (10k subset) Wikitext RULER MMLU GSM8K ARC HellaSwag OpenBookQA PIQA RACE SciQ
Benchmarks: Wikitext perplexity RULER MMLU GSM8K ARC HellaSwag OpenBookQA PIQA RACE SciQ Winogrande MQRAR
Metrics: negative log-likelihood perplexity accuracy average accuracy
Baselines: Softmax+RoPE ALiBi FIRE NoPE Gemma2-2B Qwen1.5-4B TinyLlama MQAR-style Transformers
Builds on: FlashAttention-2 RoPE ALiBi Geometric attention Llama 2
Closest papers by meaning
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention Calibration (ICML 2024 · similarity 0.89)
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.89)
- Punctuation-aware Hybrid Trainable Sparse Attention for Large Language Models (NeurIPS 2026 · similarity 0.88)
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.88)
- When Attention Sink Emerges in Language Models: An Empirical View (ICLR 2025 · similarity 0.88)
- HyperAttention: Long-context Attention in Near-Linear Time (ICLR 2024 · similarity 0.88)
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.