注意力头维度如何影响语言模型的长上下文检索
Head Dimension Matters in Language Models for Long-Context Retrieval
AI 导读
一项 arXiv 研究考察多头注意力中 head dimension 对长上下文检索的影响：在固定模型规模下，预训练损失和短上下文非检索任务几乎不受头维度影响，而长上下文检索性能随头维度变宽而提升。实验显示更宽的注意力头拥有更高比例的检索头；即使长上下文持续预训练时冻结注意力投影权重，宽头仍带来更好的检索表现，说明其还能促进 embeddings 和 MLP 等组件的有效适配。
正文 · 原文
Abstract:Multi-head attention retrieves and combines context representations through query--key interactions and value aggregation, with head dimension constraining the per-head representation space. In long-context processing, it remains unclear whether allocating a fixed parameter budget to fewer, wider heads improves information retrieval. In this paper, we investigate how head dimension affects the performance of long-context retrieval tasks. Our empirical studies reveal that pretraining loss and short-context non-retrieval task performance are not affected much by head dimension, while long-context retrieval performance improves with wider heads at fixed model sizes. We find that models with wider heads have a higher proportion of retrieval heads, associating more parameters with retrieval. Our ablation studies also show that the retrieval performance improves with wider heads even when the projection weights in attention modules are frozen during long-context continual pretraining, suggesting that wider heads not only benefit updates in attention modules but lead to more effective adaptation of components such as embeddings and MLPs.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL) |
| Cite as: | arXiv:2610.32213 [cs.LG] |
| (or arXiv:2610.32213v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.32213 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Seiya Tokui [view email]
[v1]
Fri, 23 Oct 2026 06:08:34 UTC (1,089 KB)
来源：arXiv cs.AI · arxiv.org