Specialization of softmax attention heads: insights from the high-dimensional single-location model

• URL: http://arxiv.org/abs/2603.03993v1
• Date: Wed, 04 Mar 2026 12:37:08 GMT
• Title: Specialization of softmax attention heads: insights from the high-dimensional single-location model
• Authors: M. Sagitova, O. Duranthon, L. Zdeborová,
• Abstract summary: We study the training dynamics of multi-head softmax attention under SGD.<n>We show that softmax-1 significantly reduces noise from irrelevant heads.<n>We introduce the Bayes-softmax attention, which achieves optimal prediction performance in this setting.
• Score: 0.0
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Multi-head attention enables transformer models to represent multiple attention patterns simultaneously. Empirically, head specialization emerges in distinct stages during training, while many heads remain redundant and learn similar representations. We propose a theoretical model capturing this phenomenon, based on the multi-index and single-location regression frameworks. In the first part, we analyze the training dynamics of multi-head softmax attention under SGD, revealing an initial unspecialized phase followed by a multi-stage specialization phase in which different heads sequentially align with latent signal directions. In the second part, we study the impact of attention activation functions on performance. We show that softmax-1 significantly reduces noise from irrelevant heads. Finally, we introduce the Bayes-softmax attention, which achieves optimal prediction performance in this setting.

Related papers

• SAS: Simulated Attention Score [75.1409882298863]
  We introduce Simulated Attention Score (SAS), which maintains a compact model size while simulating a larger number of attention heads and hidden feature dimension per head.<n>Comprehensive experiments on a variety of datasets and tasks demonstrate the effectiveness of the proposed SAS method.
  arXiv  Detail & Related papers  (2025-07-10T12:16:16Z)
• Beyond Parallelism: Synergistic Computational Graph Effects in Multi-Head Attention [6.670253494063657]
  Multi-head attention powers Transformer networks, the primary deep learning architecture behind the success of large language models (LLMs)<n>Yet, the theoretical advantages of multi-head versus single-head attention, beyond mere parallel processing, remain underexplored.<n>We reframe multi-head attention as a system of potentially synergistic computational graphs, where each head functions as a feedforward directed acyclic graph (DAG) with a common sink state.
  arXiv  Detail & Related papers  (2025-06-28T11:35:31Z)
• In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention [52.159541540613915]
  We study how multi-head softmax attention models are trained to perform in-context learning on linear data.<n>Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution.
  arXiv  Detail & Related papers  (2025-03-17T02:00:49Z)
• Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs [77.66717051042032]
  Practitioners have consistently observed three puzzling phenomena in transformer-based large language models. These phenomena are characterized by certain so-called "sink tokens" receiving disproportionately high attention weights. We elucidate the mechanisms behind extreme-token phenomena.
  arXiv  Detail & Related papers  (2024-10-17T17:54:06Z)
• Long-Sequence Recommendation Models Need Decoupled Embeddings [49.410906935283585]
  We identify and characterize a neglected deficiency in existing long-sequence recommendation models.<n>A single set of embeddings struggles with learning both attention and representation, leading to interference between these two processes.<n>We propose the Decoupled Attention and Representation Embeddings (DARE) model, where two distinct embedding tables are learned separately to fully decouple attention and representation.
  arXiv  Detail & Related papers  (2024-10-03T15:45:15Z)
• Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality [54.20763128054692]
  We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics.
  arXiv  Detail & Related papers  (2024-02-29T18:43:52Z)
• Superiority of Multi-Head Attention in In-Context Linear Regression [39.469021333473435]
  We conduct an exact theoretical analysis to demonstrate that multi-head attention with a substantial embedding dimension performs better than single-head attention. In general, multi-head attention is preferred over single-head attention.
  arXiv  Detail & Related papers  (2024-01-30T20:29:06Z)
• A Dynamic Head Importance Computation Mechanism for Neural Machine Translation [22.784419165117512]
  Multiple parallel attention mechanisms that use multiple attention heads facilitate greater performance of the Transformer model for various applications. In this work, we focus on designing a Dynamic Head Importance Computation Mechanism (DHICM) to dynamically calculate the importance of a head with respect to the input. We add an extra loss function to prevent the model from assigning same score to all heads, to identify more important heads and improvise performance.
  arXiv  Detail & Related papers  (2021-08-03T09:16:55Z)
• Repulsive Attention: Rethinking Multi-head Attention as Bayesian Inference [68.12511526813991]
  We provide a novel understanding of multi-head attention from a Bayesian perspective. We propose a non-parametric approach that explicitly improves the repulsiveness in multi-head attention. Experiments on various attention models and applications demonstrate that the proposed repulsive attention can improve the learned feature diversity.
  arXiv  Detail & Related papers  (2020-09-20T06:32:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.