# Selective Attention: Enhancing Transformer through Principled Context Control

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 55 references that are in the map.
- First posted on arXiv 19 Nov 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [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)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [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)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [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)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Multitask Prompted Training Enables Zero-Shot Task Generalization](https://mlsanity.com/paper/2110.08207/) (2021)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Long-tail learning via logit adjustment](https://mlsanity.com/paper/2007.07314/) (2020)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Long Range Language Modeling via Gated State Spaces](https://mlsanity.com/paper/2206.13947/) (2022)
- [Landmark Attention: Random-Access Infinite Context Length for Transformers](https://mlsanity.com/paper/2305.16300/) (2023)
- [Simple linear attention language models balance the recall-throughput tradeoff](https://mlsanity.com/paper/2402.18668/) (2024)

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

- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [Semantic XPath: Structured Agentic Memory Access for Conversational AI](https://mlsanity.com/paper/2603.01160/) (2026)
- [Attention Smoothing Is All You Need For Unlearning](https://mlsanity.com/paper/2603.01285/) (2026)
- [SpiLiFormer: Enhancing Spiking Transformers with Lateral Inhibition](https://mlsanity.com/paper/2503.15986/) (2025)
- [DSAS: A Universal Plug-and-Play Framework for Attention Optimization in Multi-Document Question Answering](https://mlsanity.com/paper/2510.12251/) (2025)
- [Scaling Limits of Long-Context Transformers](https://mlsanity.com/paper/2605.08505/) (2026)
- [OmniSparse: Training-Aware Fine-Grained Sparse Attention for Long-Video MLLMs](https://mlsanity.com/paper/2511.12201/) (2025)
- [PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference](https://mlsanity.com/paper/2503.23274/) (2025)
- [TranX-Adapter: Bridging Artifacts and Semantics within MLLMs for Robust AI-generated Image Detection](https://mlsanity.com/paper/2602.21716/) (2026)
- [Screening Is Enough](https://mlsanity.com/paper/2604.01178/) (2026)
- [Can We Go Beyond Visual Features? Neural Tissue Relation Modeling for Relational Graph Analysis in Non-Melanoma Skin Histology](https://mlsanity.com/paper/2512.06949/) (2025)
- [Attention-guided Evidence Grounding for Spoken Question Answering](https://mlsanity.com/paper/2603.16292/) (2026)
- [Good Token Hunting: A Hitchhiker's Guide to Token Selection for Visual Geometry Transformers](https://mlsanity.com/paper/2605.23892/) (2026)
- [Enhancing Multi-Agent Communication through Attention Steering with Context Relevance](https://mlsanity.com/paper/2605.30136/) (2026)
- [Pre-registered tests of solid-state-physics-inspired LLM compression: a cluster-level negative result at small-language-model scale](https://mlsanity.com/paper/2609.34292/) (2026)
- [Message Passing Does More with Less for In-Context Learning on Graphs](https://mlsanity.com/paper/2609.37057/) (2026)

## Abstract (arXiv)

The attention mechanism within the transformer architecture enables the model to weigh and combine tokens based on their relevance to the query. While self-attention has enjoyed major success, it notably treats all queries q in the same way by applying the mapping V^topsoftmax(Kq), where V,K are the value and key embeddings respectively. In this work, we argue that this uniform treatment hinders the ability to control contextual sparsity and relevance. As a solution, we introduce the textitSelective Self-Attention (SSA) layer that augments the softmax nonlinearity with a principled temperature scaling strategy. By controlling temperature, SSA adapts the contextual sparsity of the attention map to the query embedding and its position in the context window. Through theory and experiments, we demonstrate that this alleviates attention dilution, aids the optimization process, and enhances the model's ability to control softmax spikiness of individual queries. We also incorporate temperature scaling for value embeddings and show that it boosts the model's ability to suppress irrelevant/noisy tokens. Notably, SSA is a lightweight method which introduces less than 0.5% new parameters through a weight-sharing strategy and can be fine-tuned on existing LLMs. Extensive empirical evaluations demonstrate that SSA-equipped models achieve a noticeable and consistent accuracy improvement on language modeling benchmarks.
