Research
Choosing the Number of Attention Heads from Intrinsic Dimension: A Small-Scale Study
arXiv:2610.19966v1 Announce Type: new 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: estima
arXiv:2610.19966v1 Announce Type: new 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.
Related
- Tracing the complexity profiles of different linguistic phenomena through the intrinsic dimension of LLM representations
- Do Syntax Heads Influence Model Behavior?
- Influence Score and Transformers interpretability: Measure of the Effective Impact of Attention Heads at inference time
- Multiscale POD of Transformer Attention Fields: Scale-Selective Analysis via Morlet Scalogram
Source: arXiv cs.CL | 2026-10-15