Research map / Architectures and efficiency / Attention and transformers
StableMask: Refining Causal Masking in Decoder-only Transformer
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 StableMask, a parameter-free modification of the causal mask that adds decaying pseudo-attention scores to address disproportional attention and restore absolute position encoding in decoder-only Transformers with RPE.
Problem
Decoder-only Transformers with causal masking and relative position encoding force attention to sum to one over all visible tokens, causing disproportional attention to initial/punctuation tokens, and they cannot encode absolute positional information well.
Approach
Append pseudo-attention scores with linearly decaying positional values (-(j-1)gamma) to the upper triangle before softmax, then zero out the upper triangle after softmax; this lets rows sum to less than one, letting heads dump excess attention and encoding absolute position, with no new parameters.
Key results
- On MiniPile, OpenLLaMA 160M perplexity improves from 25.9 to 25.0 (1 epoch) and 21.2 to 20.9 (2 epochs) with StableMask.
- On WikiText-103, OpenLLaMA 160M perplexity drops from 22.5 to 21.1 with StableMask.
- On synthetic absolute-position tasks, RPE models (ALiBi, T5, RoPE) reach only about 21-25% accuracy on task (1) versus 96-98% for APE models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add StableMask | BLOOM 71M with ALiBi · WikiText-103 | PPL: 29.9 -> 29.0 | Helped |
| Add StableMask | OpenLLaMA 71M with RoPE · WikiText-103 | PPL: 27.4 -> 26.9 | Helped |
| Add StableMask | BLOOM 160M with ALiBi · WikiText-103 | PPL: 27.6 -> 26.1 | Helped |
| Add StableMask | OpenLLaMA 160M with RoPE · WikiText-103 | PPL: 22.5 -> 21.1 | Helped |
| Add StableMask | BLOOM 160M with ALiBi · MiniPile, 1 epoch | PPL: 25.8 -> 25.6 | Helped |
| Add StableMask | OpenLLaMA 160M with RoPE · MiniPile, 1 epoch | PPL: 25.9 -> 25.0 | Helped |
| Add StableMask | BLOOM 160M with ALiBi · MiniPile, 2 epochs | PPL: 23.3 -> 22.9 | Helped |
| Add StableMask | OpenLLaMA 160M with RoPE · MiniPile, 2 epochs | PPL: 21.2 -> 20.9 | Helped |
| Add StableMask | BLOOM 430M with ALiBi · MiniPile | PPL: 20.6 -> 19.6 (1 epoch); 15.6 -> 15.5 (2 epochs) | Helped |
| Add StableMask | OpenLLaMA 430M with RoPE · MiniPile | PPL: 19.6 -> 19.5 (1 epoch); 15.7 -> 15.1 (2 epochs) | Helped |
| Add StableMask | OpenLLaMA 1.4B with RoPE · Pile, 5B–25B tokens | PPL: 15.4 -> 15.0; 14.8 -> 14.6; 12.4 -> 11.9; 11.7 -> 11.3; 10.7 -> 10.4 | Helped |
| Add zero-valued pseudo-attention scores (ablation) | StableMask on OpenLLaMA 160M · MiniPile | PPL: 22.5 -> 21.5 | Helped |
Surprising findings
- Standard RPE Transformers fail badly on simple absolute-position tasks despite strong language modeling performance.
Limitations
- Original StableMask cannot use KV caching at inference, requiring a suffix-token workaround.
- Slightly increased computation compared to standard attention.
- Absolute position encoding requires careful calibration to avoid harming the model.
Experimental setup
Datasets: WikiText-103 MiniPile The Pile LAMBADA PIQA ARC-Easy ARC-Challenge OpenbookQA Winogrande
Benchmarks: WikiText-103 MiniPile The Pile LAMBADA PIQA ARC-Easy ARC-Challenge OpenbookQA Winogrande
Metrics: perplexity accuracy
Baselines: OpenLLaMA BLOOM Learnable AT Fixed Value AT URPE StreamingLLM
Builds on: RoPE ALiBi FlashAttention URPE StreamingLLM
Closest papers by meaning
- Reconsidering Positional Supervision in Masked Diffusion Language Model Training (NeurIPS 2026 · similarity 0.88)
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.87)
- PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels (ICML 2024 · similarity 0.87)
- Memory-Augmenting Decoder-Only Language Models through Encoders (Student Abstract) (AAAI 2024 · similarity 0.87)
- Functional Interpolation for Relative Positions improves Long Context Transformers (ICLR 2024 · similarity 0.86)
- Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding (ICML 2025 · similarity 0.86)
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.