# StableMask: Refining Causal Masking in Decoder-only Transformer

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 35 references that are in the map.
- First posted on arXiv 7 Feb 2024.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [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)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Rethinking Positional Encoding in Language Pre-training](https://mlsanity.com/paper/2006.15595/) (2020)
- [Transformer Quality in Linear Time](https://mlsanity.com/paper/2202.10447/) (2022)
- [Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing](https://mlsanity.com/paper/2306.12929/) (2023)
- [The MiniPile Challenge for Data-Efficient Language Models](https://mlsanity.com/paper/2304.08442/) (2023)
- [Your Transformer May Not be as Powerful as You Expect](https://mlsanity.com/paper/2205.13401/) (2022)
- [Bidirectional Language Models Are Also Few-shot Learners](https://mlsanity.com/paper/2209.14500/) (2022)
- [Rethinking Softmax Cross-Entropy Loss for Adversarial Robustness](https://mlsanity.com/paper/1905.10626/) (2019)

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

- [When Attention Sink Emerges in Language Models: An Empirical View](https://mlsanity.com/paper/2410.10781/) (2024)
- [Seeing Far and Clearly: Mitigating Hallucinations in MLLMs with Attention Causal Decoding](https://mlsanity.com/paper/2505.16652/) (2025)
- [Deeper Insights Without Updates: The Power of In-Context Learning Over Fine-Tuning](https://mlsanity.com/paper/2410.04691/) (2024)
- [When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context Training](https://mlsanity.com/paper/2411.13476/) (2024)
- [MiSS: Revisiting the Trade-off in LoRA with an Efficient Shard-Sharing Structure](https://mlsanity.com/paper/2409.15371/) (2024)
- [Rethinking Causal Mask Attention for Vision-Language Inference](https://mlsanity.com/paper/2505.18605/) (2025)
- [A Structural Theory of Position Bias in Transformers](https://mlsanity.com/paper/2602.16837/) (2026)
- [AERO: Entropy-Guided Framework for Private LLM Inference](https://mlsanity.com/paper/2410.13060/) (2024)
- [DenoMAE2.0: Improving Denoising Masked Autoencoders by Classifying Local Patches](https://mlsanity.com/paper/2502.18202/) (2025)
- [Induction Head Toxicity Mechanistically Explains Repetition Curse in Large Language Models](https://mlsanity.com/paper/2505.13514/) (2025)
- [Aligning What Vision-Language Models See and Perceive with Adaptive Information Flow](https://mlsanity.com/paper/2604.15809/) (2026)
- [NanoVLMs: How small can we go and still make coherent Vision Language Models?](https://mlsanity.com/paper/2502.07838/) (2025)
- [MODA: MOdular Duplex Attention for Multimodal Perception, Cognition, and Emotion Understanding](https://mlsanity.com/paper/2507.04635/) (2025)
- [Rethinking Key-frame-based Micro-expression Recognition: A Robust and Accurate Framework Against Key-frame Errors](https://mlsanity.com/paper/2508.06640/) (2025)
- [Making Transformers Logical](https://mlsanity.com/paper/2608.00418/) (2026)
- [CaTs and DAGs: Integrating Directed Acyclic Graphs with Transformers for Causally Constrained Predictions](https://mlsanity.com/paper/2410.14485/) (2024)
- [ModRWKV: Transformer Multimodality in Linear Time](https://mlsanity.com/paper/2505.14505/) (2025)
- [Enhancing Instruction-Following Capabilities in Seq2Seq Models: DoLA Adaptations for T5](https://mlsanity.com/paper/2512.03803/) (2025)
- [A Circuit, Not The Circuit: Non-Unique Causal Localisation of the Mamba-2 State Sink](https://mlsanity.com/paper/2606.00930/) (2026)
- [From Attention Masks to Inert Zero-Vector Tokens: OAttention and O-Closure for Token Dynamics](https://mlsanity.com/paper/2608.21174/) (2026)

## Abstract (arXiv)

The decoder-only Transformer architecture with causal masking and relative position encoding (RPE) has become the de facto choice in language modeling. Despite its exceptional performance across various tasks, we have identified two limitations: First, it requires all attention scores to be non-zero and sum up to 1, even if the current embedding has sufficient self-contained information. This compels the model to assign disproportional excessive attention to specific tokens. Second, RPE-based Transformers are not universal approximators due to their limited capacity at encoding absolute positional information, which limits their application in position-critical tasks. In this work, we propose StableMask: a parameter-free method to address both limitations by refining the causal mask. It introduces pseudo-attention values to balance attention distributions and encodes absolute positional information via a progressively decreasing mask ratio. StableMask's effectiveness is validated both theoretically and empirically, showing significant enhancements in language models with parameter sizes ranging from 71M to 1.4B across diverse datasets and encoding methods. We further show that it naturally supports (1) efficient extrapolation without special tricks such as StreamingLLM and (2) easy integration with existing attention optimization techniques.