arXiv cs.CL· Seiya Tokui, Shoichiro Yamaguchi, Sosuke Kobayashi·· 12 小时前AI 评分31
语言模型长上下文检索中注意力头维度的影响研究
Head Dimension Matters in Language Models for Long-Context Retrieval
AI 导读
在固定模型规模下，注意力头维度更宽可提升长上下文检索性能，而预训练损失和短上下文非检索任务表现基本不受影响。研究发现宽头模型的检索头占比更高，把更多参数分配给检索；即使长上下文持续预训练时冻结注意力模块的投影权重，宽头仍能带来检索提升，说明其还促进了嵌入向量和 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.CL · arxiv.org