语言模型中用于长上下文检索的注意力头维度至关重要
Head Dimension Matters in Language Models for Long-Context Retrieval
精选理由
揭示增大注意力头维度可提升长上下文检索能力，为模型架构设计提供依据。
AI 摘要
arXiv:2610.32213v1 发布类型：新 摘要：多头注意力通过查询--键交互和数值聚合来检索并组合上下文表示，其中头维度限制了每个注意力头的表示空间。在长上下文处理中，在固定参数预算下分配给更少但更宽的注意力头是否能改善信息检索，仍不明确。本文研究了头维度如何影响长上下文检索任务的性能。我们的实证研究表明，预训练损失和短上下文非检索任务的性能受头维度影响不大，而在模型规模固定的情况下，更宽的注意力头能够提升长上下文检索性能。我们发现，具有更宽注意力头的模型拥有更高比例的检索头，将更多参数用于检索。我们的消融研究还表明，即使在长上下文持续预训练期间冻结注意力模块中的投影权重，更宽的注意力头仍能提升检索性能，这表明更宽的注意力头不仅有利于注意力模块中的参数更新，还能促使嵌入层和 MLP 等组件实现更有效的适应。