旧金山 – 在自然语言处理领域，Transformer模型及其注意力机制一直是研究的热点。然而，当处理包含大量Token的上下文时，标准注意力机制往往难以精准捕捉关键信息，忽略干扰因素。近日，Meta公司发布了一种名为Multi-Token Attention (MTA) 的新型Transformer注意力机制，旨在突破传统注意力机制的瓶颈，提升模型在复杂上下文中的表现。
标准注意力的局限性
标准多头注意力机制通过计算查询向量与上下文Token对应的键向量之间的点积相似性来确定注意力权重。与查询相似的关键字会获得更高的权重，从而影响输出向量。例如，当查询向量与“Alice”Token相关时，模型可以定位上下文中所有提及“Alice”的内容。然而，这种机制的根本限制在于，每个注意力权重仅取决于单个关键字和查询向量，这使得模型难以识别需要多个Token共同表达的上下文信息。
例如，要查找同时提到“Alice”和“rabbit”的句子，查询向量需要同时编码这两个Token的信息。虽然可以通过Transformer的层将多个Token编码成一个向量，但这需要增加模型维度，并消耗大量计算资源。
Multi-Token Attention的创新
为了解决上述问题，Meta的研究人员提出了Multi-Token Attention (MTA)，其核心思想是利用多个向量对的相似性来确定注意力集中的位置。MTA通过对现有注意力机制进行简单的修改来实现这一目标，即在注意力权重上进行卷积运算，该运算在键、查询和注意力头三个维度上运行。这种设计允许注意力权重以相邻键、之前的查询和其他头为条件。
具体来说，MTA可以先分别查找“Alice”和“rabbit”的提及，然后将这些注意力组合在一起，只关注两者同时存在的地方。
MTA的技术细节
MTA由三个关键部分组成：
- 键-查询卷积（key-query convolution）： 在注意力logit上进行卷积操作，结合来自多个查询和键Token的信息。
- 头混合卷积（head mixing convolution）： 在头组中使用头卷积，将不同头的注意力权重组合起来，实现跨头信息共享。
- 带深度缩放的组归一化： 抵消残差流，改善梯度流。
实验结果与结论
研究人员通过一系列实验验证了MTA的有效性。首先，在一个玩具任务中，MTA轻松解决了标准注意力机制难以处理的问题。其次，在一个包含1050亿个词库的880M参数模型上进行预训练，结果表明MTA在验证复杂度和标准基准任务方面都有所改进，而参数数量仅增加了0.001%。此外，在长语境任务（如Needle-in-the-Haystack和BabiLong）上，MTA的表现也明显优于基线模型。
未来展望
Multi-Token Attention的提出为Transformer模型的注意力机制带来了新的思路，有望在自然语言处理领域取得更广泛的应用。该研究表明，通过对现有注意力机制进行巧妙的改进，可以有效提升模型在复杂上下文中的表现。未来，我们可以期待更多基于MTA的创新应用，例如在机器翻译、文本摘要和对话生成等领域。
参考文献
- 论文：Multi-Token Attention
- 论文链接：https://arxiv.org/abs/2504.00927
（本文部分信息来源于机器之心报道）
Views: 0