# Affine-Scaled Attention: Towards Flexible and Stable Transformer Attention

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 26 references that are in the map.
- First posted on arXiv 26 Feb 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [When Attention Sink Emerges in Language Models: An Empirical View](https://mlsanity.com/paper/2410.10781/) (2024)
- [A Comprehensive Survey on Applications of Transformers for Deep Learning Tasks](https://mlsanity.com/paper/2306.07303/) (2023)
- [Bridging the Divide: Reconsidering Softmax and Linear Attention](https://mlsanity.com/paper/2412.06590/) (2024)
- [Approximation Rate of the Transformer Architecture for Sequence Modeling](https://mlsanity.com/paper/2305.18475/) (2023)
- [What are you sinking? A geometric approach on attention sink](https://mlsanity.com/paper/2508.02546/) (2025)
- [Understanding the Expressive Power and Mechanisms of Transformer for Sequence Modeling](https://mlsanity.com/paper/2402.00522/) (2024)
- [In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness](https://mlsanity.com/paper/2402.11639/) (2024)
- [Normalized Attention Without Probability Cage](https://mlsanity.com/paper/2005.09561/) (2020)
- [Why Softmax Attention Outperforms Linear Attention](https://mlsanity.com/paper/2310.11685/) (2023)
- [On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective](https://mlsanity.com/paper/2507.23632/) (2025)

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

- [OmniOpt: Taxonomy, Geometry, and Benchmarking of Modern Optimizers](https://mlsanity.com/paper/2607.04033/) (2026)
- [What Softmax Throws Away: Mass-Aware Attention for Evidence Accumulation](https://mlsanity.com/paper/2607.22781/) (2026)
- [Abstention and Noise Filtering: Two Missing Primitives of Softmax Attention](https://mlsanity.com/paper/2609.22005/) (2026)

## Abstract (arXiv)

Transformer attention is typically implemented using softmax normalization, which enforces attention weights with unit sum normalization. While effective in many settings, this constraint can limit flexibility in controlling attention magnitudes and may contribute to overly concentrated or unstable attention patterns during training. Prior work has explored modifications such as attention sinks or gating mechanisms, but these approaches provide only limited or indirect control over attention reweighting. We propose Affine-Scaled Attention, a simple extension to standard attention that introduces input-dependent scaling and a corresponding bias term applied to softmax-normalized attention weights. This design relaxes the strict normalization constraint while maintaining aggregation of value representations, allowing the model to adjust both the relative distribution and the scale of attention in a controlled manner. We empirically evaluate Affine-Scaled Attention in large-scale language model pretraining across multiple model sizes. Experimental results show consistent improvements in training stability, optimization behavior, and downstream task performance compared to standard softmax attention and attention sink baselines. These findings suggest that modest reweighting of attention outputs provides a practical and effective way to improve attention behavior in Transformer models.
