Choosing the Number of Attention Heads from Intrinsic Dimension: A Small-Scale Study

Arjun Boya; Charan Sai Bandaru; Masthanaiah Pillimetla; Kavana Chidanandappa

cs.CLOct 14, 2026

Choosing the Number of Attention Heads from Intrinsic Dimension: A Small-Scale Study

Authors: Arjun Boya, Charan Sai Bandaru, Masthanaiah Pillimetla, Kavana Chidanandappa

Organizations: Institute of Cognitive Science University of Osnabrück · Department of Social Sciences and Technology Technical University of Munich

Abstract

The number of attention heads in a transformer is usually set by convention or grid search. We test a simple heuristic for setting it in advance: estimate the intrinsic dimension of token representations without labels, and use roughly that many heads. We estimate intrinsic dimension with a nested-dropout VAE trained on DistilBERT embeddings of WikiText-2, choosing the KL weight at the knee of the VAE's rate-distortion curve rather than by hand. This gives $k^*=9.2$ (95% CI [8.6, 9.8]). In a small 4-layer transformer trained from scratch, 10 heads and 32 heads are statistically equivalent on SST-2 and on MRPC (paired TOST over 10 seeds, $\pm1$ point margin), with 69% fewer attention parameters for the 10-head model. Training 10 heads from the start also slightly outperforms training 32 heads and pruning to 10 (0.6 points, $p=0.008$ ). The evidence is narrow: two classification tasks, one small architecture, no pretrained models, and no formal link between intrinsic dimension and head count. We present the result as a practical heuristic worth testing at larger scale, not as a general principle.

Figures & tables

Panel of Figure 1 — Figure 1: The $\beta$ sweep (mean over 5 seeds). Left: held-out reconstruction loss; $\beta^{*}$ is the point furthest from the dashed chord. Right: elbow of the latent variance spectrum, with 95% confidence intervals. At $\beta^{*}=4$ the elbow is $k^{*}=9.2$ ; the dotted line marks the head count tested downstream.

$\beta$	Reconstruction loss	Elbow $k$
0.10	0.354	$34.0\pm 0.9$
0.25	0.382	$26.4\pm 1.1$
0.50	0.422	$21.6\pm 1.4$
1.00	0.485	$17.0\pm 1.5$
2.00	0.602	$13.0\pm 0.9$
3.00	0.707	$10.4\pm 0.7$

Table 1: Values behind Figure 1 : held-out reconstruction loss and elbow

k

(mean

\pm

95% CI half-width over 5 seeds). Italic rows were added after the first seven-point sweep.

Task	$k=10$	$k=32$	Paired $t$ -test	TOST
SST-2	$82.13\%\pm 0.42$	$81.95\%\pm 0.33$	$p=0.31$	equivalent, $p=0.0005$
MRPC	$70.47\%\pm 0.90$	$70.37\%\pm 1.03$	$p=0.74$	equivalent, $p=0.006$

Table 2: Validation accuracy of 10 versus 32 heads (mean

\pm

std over 10 paired seeds). TOST tests equivalence within

\pm 1

accuracy point.

Configuration	Validation accuracy
32 heads, unpruned	$81.95\%\pm 0.33$
32 pruned to 10, before fine-tuning	$71.69\%\pm 7.04$
32 pruned to 10, after fine-tuning	$81.50\%\pm 0.37$
10 heads from the start	$\mathbf{82.13\%\pm 0.42}$

Table 3: Pruning baseline on SST-2 (mean

\pm

std over the same 10 seeds).

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

	SST-2	MRPC	Pruned fine-tune	VAE
Learning rate	$10^{-4}$	$2\times 10^{-4}$	$2\times 10^{-5}$	$10^{-3}$
Weight decay	0.05	0.03	0.05	$10^{-4}$
Dropout	0.30	0.25	0.30	–
Batch size	64	32	64	256
Max epochs	10	20	5	40
Early-stopping patience	3	6	2	–

Appendix

Table 4: Training settings. All classifier runs use AdamW, a cosine learning-rate schedule with linear warmup, gradient clipping at 1.0, and early stopping on validation accuracy.

Multi-Head Attention

Multi-head attention uses parallel attention heads to route and combine information in Transformers, shaping their capacity, interpretability, and efficiency.

See all 73 papers

Multi-Head Attention as Ensemble Nadaraya-Watson Estimation: Variance Reduction, Decorrelation, and Optimal Head Diversity
May 18, 2026Ernest FokouéSelf-Attention Ensemble Learning
School of Mathematics and Statistics, College of Science Rochester Institute of Technology Rochester, New York 14623, USA
Performance-Efficiency Tradeoffs in Transformers: An Approximation Theory Perspective
Oct 4, 2025Ruoxi Yu, Haotian Jiang, Jingpu Cheng +3Softmax Attention Attention Head Analysis
Center for Data Science, Peking University · Institute for Functional Intelligent Materials, National University of Singapore · Department of Mathematics, National University of Singapore +1
INCRT: An Incremental Transformer That Determines Its Own Architecture
Apr 12, 2026Giansalvo CirrincioneTransformer Neural Architecture Search
Laboratoire LTI, Université de Picardie Jules Verne, Amiens, France

May 18, 2026J/K move · Enter open · S save

Multi-Head Attention as Ensemble Nadaraya-Watson Estimation: Variance Reduction, Decorrelation, and Optimal Head Diversity

Ernest Fokoué

School of Mathematics and Statistics, College of Science Rochester Institute of Technology Rochester, New York 14623, USA

Self-Attention Ensemble Learning Transformer Attention Statistical Learning Theory Multi-Head Attention Variance Reduction+1

Open paper AskPDF arXiv

We develop a rigorous statistical theory of multi-head attention (MHA) as an ensemble of Nadaraya-Watson (NW) kernel regression estimators. Building on the algebraic identity between single-head softmax attention and the NW estimator, we prove that MHA is a structured ensemble of H NW estimators, each operating in a distinct learned projection subspace of the key space. We derive an explicit Bias-Variance-Covariance decomposition of the MHA mean squared error, showing that variance reduction depends not merely on the number of heads H but fundamentally on the decorrelation of head outputs. Decorrelation is governed by the principal angles between learned projection subspaces: orthogonal projections yield maximum variance reduction; aligned projections yield none. We introduce the Head Diversity Index (HDI), a computable spectral measure of inter-head decorrelation, and prove that MHA mean squared error is monotonically decreasing in HDI. This provides the first rigorous theoretical explanation for the empirically observed specialization of attention heads. Under a fixed total-dimension budget D = H * d_k, we solve the optimal head-dimension allocation problem, deriving the MSE-minimizing pair (H*, d_k*) from data distribution and regression smoothness. The solution yields a new architectural scaling law: the optimal per-head dimension grows logarithmically with training set size, while the optimal number of heads grows nearly linearly with the total budget D. Our framework unifies three strands of prior work: the NW theory of single-head attention, the general weighting theory for ensemble learning, and the decorrelation-variance-reduction isomorphism between biological and computational ensembles. Multi-head attention is the Transformer's instantiation of a universal principle: identical agents plus diversity-enforcing mechanisms yields emergent optimality.

Abstract

Multi-Head Attention

Explore similar work

Multi-Head Attention as Ensemble Nadaraya-Watson Estimation: Variance Reduction, Decorrelation, and Optimal Head Diversity

Performance-Efficiency Tradeoffs in Transformers: An Approximation Theory Perspective

INCRT: An Incremental Transformer That Determines Its Own Architecture