arXiv:2504.00927cs.CL2025-04被引 12

让大模型同时参考多个关键词,提升长文本定位能力

Multi-Token Attention

  • 用卷积融合多个查询与键向量,实现多标记注意力
  • 在长上下文任务中超越标准Transformer模型表现
  • 适合需要精准检索长篇内容的场景

软注意力是大语言模型定位上下文相关部分的关键机制。然而,传统方法仅基于单个查询与键向量的相似性计算注意力权重,这种‘单标记注意力’限制了信息利用效率。为此,我们提出多标记注意力(Multi-Token Attention, MTA),使模型能同时基于多个查询和键向量确定注意力权重。通过在查询、键和注意力头维度上引入卷积操作,使邻近的查询与键相互影响注意力分布,从而实现更精确的定位。实验表明,该方法在多种主流基准测试中均取得更好性能,尤其在长上下文信息检索任务中优势明显,显著优于标准Transformer基线模型。

原文摘要 · Abstract (English)

Soft attention is a critical mechanism powering LLMs to locate relevant parts within a given context. However, individual attention weights are determined by the similarity of only a single query and key token vector. This "single token attention" bottlenecks the amount of information used in distinguishing a relevant part from the rest of the context. To address this issue, we propose a new attention method, Multi-Token Attention (MTA), which allows LLMs to condition their attention weights on multiple query and key vectors simultaneously. This is achieved by applying convolution operations over queries, keys and heads, allowing nearby queries and keys to affect each other's attention weights for more precise attention. As a result, our method can locate relevant context using richer, more nuanced information that can exceed a single vector's capacity. Through extensive evaluations, we demonstrate that MTA achieves enhanced performance on a range of popular benchmarks. Notably, it outperforms Transformer baseline models on standard language modeling tasks, and on tasks that require searching for information within long contexts, where our method's ability to leverage richer information proves particularly beneficial.

注意力机制大模型长文本

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。