The Effect of Attention Head Count on Transformer Approximation

View PDF HTML (experimental)

Abstract:Transformer has become the dominant architecture for sequence modeling, yet a detailed understanding of how its structural parameters influence expressive power remains limited. In this work, we study the approximation properties of transformers, with particular emphasis on the role of the number of attention heads. Our analysis begins with the introduction of a generalized $D$-retrieval task, which we prove to be dense in the space of continuous functions, thereby providing the basis for our theoretical framework. We then establish both upper and lower bounds on the parameter complexity required for $\epsilon$-approximation. Specifically, we show that transformers with sufficiently many heads admit efficient approximation, whereas with too few heads, the number of parameters must scale at least as $O(1/\epsilon^{cT})$, for some constant $c$ and sequence length $T$. To the best of our knowledge, this constitutes the first rigorous lower bound of this type in a nonlinear and practically relevant setting. We further examine the single-head case and demonstrate that an embedding dimension of order $O(T)$ allows complete memorization of the input, where approximation is entirely achieved by the feed-forward block. Finally, we validate our theoretical findings with experiments on both synthetic data and real-world tasks, illustrating the practical relevance of our results.

Comments: Accepted by ICLR 2026

Subjects:

Machine Learning (cs.LG); Machine Learning (stat.ML)

Cite as: arXiv:2510.06662 [cs.LG]

(or arXiv:2510.06662v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2510.06662

arXiv-issued DOI via DataCite

Submission history

From: Penghao Yu [view email] [v1] Wed, 8 Oct 2025 05:27:25 UTC (443 KB) [v2] Tue, 31 Mar 2026 07:14:19 UTC (454 KB)