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, ±1 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
Figure 1: The β sweep (mean over 5 seeds). Left: held-out reconstruction loss; β∗ is the point furthest from the dashed chord. Right: elbow of the latent variance spectrum, with 95% confidence intervals. At β∗=4 the elbow is k∗=9.2 ; the dotted line marks the head count tested downstream.
β
Reconstruction loss
Elbow k
0.10
0.354
34.0±0.9
0.25
0.382
26.4±1.1
0.50
0.422
21.6±1.4
1.00
0.485
17.0±1.5
2.00
0.602
13.0±0.9
3.00
0.707
10.4±0.7
Table 1: Values behind Figure 1 : held-out reconstruction loss and elbow k (mean ± 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%±0.42
81.95%±0.33
p=0.31
equivalent, p=0.0005
MRPC
70.47%±0.90
70.37%±1.03
p=0.74
equivalent, p=0.006
Table 2: Validation accuracy of 10 versus 32 heads (mean ± std over 10 paired seeds). TOST tests equivalence within ±1 accuracy point.
Configuration
Validation accuracy
32 heads, unpruned
81.95%±0.33
32 pruned to 10, before fine-tuning
71.69%±7.04
32 pruned to 10, after fine-tuning
81.50%±0.37
10 heads from the start
82.13%±0.42
Table 3: Pruning baseline on SST-2 (mean ± 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×10−4
2×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 uses parallel attention heads to route and combine information in Transformers, shaping their capacity, interpretability, and efficiency.
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.
Ernest Fokoué
School of Mathematics and Statistics, College of Science Rochester Institute of Technology Rochester, New York 14623, USA
Transformers have achieved remarkable successes across a wide range of applications, yet the theoretical foundation of their model efficiency remains underexplored. In this work, we investigate how the model parameters -- mainly attention heads and head dimensions -- should be allocated across layers to balance expressivity and efficiency. We first provide mathematical analysis on the role of early layers in information extraction from an approximation perspective, with a theoretical characterization on the trade-off between the number of heads and head dimension under a fixed parameter budget. In addition, we uncover and prove the saturation behavior of softmax activations: Continuously increasing head dimensions can lead to diminishing returns in learning errors, particularly for long sequences. Supported by both theory and experiments, this saturation pattern suggests that later layers can operate more efficiently with reduced parameters. Combining these insights, we propose principled strategies for allocating attention heads and dimensions across Transformers' layers, shedding light on theoretically-grounded model efficiency of Transformer-based architectures.
Ruoxi Yu, Haotian Jiang, Jingpu Cheng +3
Center for Data Science, Peking University · Institute for Functional Intelligent Materials, National University of Singapore · Department of Mathematics, National University of Singapore +1
Transformer architectures are designed by trial and error: the number of attention heads, the depth, and the head size are fixed before training begins, with no mathematical principle to guide the choice. The result is systematic structural redundancy -- between half and four-fifths of all heads in a trained model can be removed without measurable loss -- because the architecture allocates capacity without reference to the actual requirements of the task.This paper introduces INCRT (Incremental Transformer), an architecture that determines its own structure during training. Starting from a single head, INCRT adds one attention head at a time whenever its current configuration is provably insufficient, and prunes heads that have become redundant. Each growth decision is driven by a single, online-computable geometric quantity derived from the task's directional structure, requiring no separate validation phase and no hand-tuned schedule. Two theorems form the theoretical backbone. The first (homeostatic convergence) establishes that the system always reaches a finite stopping configuration that is simultaneously minimal (no redundant heads) and sufficient (no uncaptured directional energy above the threshold). The second (compressed-sensing analogy) provides a geometric upper bound on the number of heads that this configuration can contain, as a function of the spectral complexity of the task. Experiments on SARS-CoV-2 variant classification and SST-2 sentiment analysis confirm both results: the predicted and observed head counts agree within 12% across all benchmarks, and the final architectures match or exceed BERT-base on distribution-specific tasks while using between three and seven times fewer parameters and no pre-training.
Giansalvo Cirrincione
Laboratoire LTI, Université de Picardie Jules Verne, Amiens, France