arXiv：cs.LG· Arjun Boya, Charan Sai Bandaru, Masthanaiah Pillimetla, Kavana Chidanandappa·· 6 小时前
从内在维度确定注意力头数量：一项小规模研究
Choosing the Number of Attention Heads from Intrinsic Dimension: A Small-Scale Study
AI 导读
一项小规模研究提出用无标签估计 token 表示的内在维度来预设 Transformer 注意力头数：在 WikiText-2 的 DistilBERT 嵌入上训练嵌套 dropout VAE，取 rate-distortion 曲线拐点处的 KL 权重，得到 k*=9.2（95% CI [8.6, 9.8]）。
正文
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.
| Comments: | 9 pages, 1 figure, 4 tables. Code: this https URL |
| Subjects: | Computation and Language (cs.CL); Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.19966 [cs.CL] |
| (or arXiv:2610.19966v1 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2610.19966 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Arjun Boya [view email]
[v1]
Wed, 14 Oct 2026 10:50:40 UTC (12 KB)
来源：arXiv：cs.LG · arxiv.org