麻省理工学院团队推出PaTH注意力机制提升大语言模型能力

大多数语言都依赖词汇位置和句子结构来传达含义。例如，"猫坐在盒子上"与"盒子在猫身上"的含义完全不同。在长篇文本中，如金融文档或小说，这些词汇的语法结构很可能会发生变化。

同样，一个人可能在代码中跟踪变量或遵循带有条件动作的指令。这些都是状态变化和序列推理的例子，我们期望最先进的人工智能系统能够在这些方面表现出色。然而，现有的Transformer前沿注意力机制——大语言模型中用于确定词汇重要性的主要架构——在这些能力方面存在理论和实证局限性。

注意力机制允许大语言模型回顾查询或文档的早期部分，并根据其训练确定哪些细节和词汇最重要。然而，这种机制本身并不理解词序。它同时"看到"所有输入词汇（也称为Token），并按照呈现的顺序处理它们，因此研究人员开发了编码位置信息的技术。这对于像语言这样高度结构化的领域至关重要。但主要的位置编码方法，称为旋转位置编码（RoPE），只考虑序列中Token之间的相对距离，与输入数据无关。这意味着，例如，相距四个位置的词汇，如上例中的"猫"和"盒子"，都会接收到针对该相对距离的相同固定数学旋转。

现在，由麻省理工学院和MIT-IBM Watson人工智能实验室领导的研究产生了一种称为"PaTH注意力"的编码技术，使位置信息变得自适应和上下文感知，而不像RoPE那样静态。

该论文的资深作者、电气工程与计算机科学系副教授、计算机科学与人工智能实验室成员、MIT-IBM Watson人工智能实验室研究员Yoon Kim说："Transformer能够准确且可扩展地建模许多领域，但它们在状态跟踪方面存在局限性，而状态跟踪是我们希望人工智能系统具备的重要能力的基础现象。因此，重要的问题是：我们如何在保持Transformer可扩展性和效率的同时，实现状态跟踪？"

本月早些时候，关于这项工作的新论文在神经信息处理系统会议（NeurIPS）上发表。Kim的合作者包括主要作者、电气工程与计算机科学系研究生、前MIT-IBM Watson人工智能实验室暑期项目实习生杨松林；斯坦福大学的温凯悦；微软的任礼良；以及IBM研究院和MIT-IBM Watson人工智能实验室的沈奕康、Shawn Tan、Mayank Mishra和Rameswar Panda。

理解路径

PaTH注意力不像RoPE那样根据Token之间的相对距离为每个词分配固定旋转，而是灵活的，将中间词汇视为由小的、数据相关转换组成的路径。每个转换基于称为Householder反射的数学运算，就像一个微小的镜子，根据它经过的每个Token的内容进行调整。序列中的每一步都可以影响模型后续如何解释信息。累积效应让系统能够建模词汇之间路径上含义如何变化，而不仅仅是它们相距多远。这种方法允许Transformer跟踪实体和关系如何随时间变化，赋予其"位置记忆"的感觉。可以将此想象为在体验环境及其对你的影响的同时行走路径。此外，团队还开发了一种硬件高效算法，更有效地计算每对Token之间的注意力分数，使得PaTH注意力的累积数学转换被压缩并分解为更小的计算，从而与GPU上的快速处理兼容。

MIT-IBM研究人员随后探索了PaTH注意力在合成和真实任务上的表现，包括推理、长上下文基准测试和完整大语言模型训练，以查看是否改善了模型跟踪信息随时间变化的能力。团队测试了其在许多干扰步骤和多步回忆测试中遵循最新"写"命令的能力，这些任务对于像RoPE这样的标准位置编码方法来说是困难的。研究人员还训练了中等规模的大语言模型并与其他方法进行比较。PaTH注意力改善了困惑度，并在未经训练的推理基准上超越了其他方法。他们还评估了使用数万个Token输入的检索、推理和稳定性。PaTH注意力始终证明了内容感知能力。

Kim说："我们发现，无论是在设计用于测试Transformer局限性的诊断任务上，还是在真实世界的语言建模任务上，我们的新方法都能够超越现有的注意力机制，同时保持其效率。"此外，"我很期待看到这些类似PATH的数据相关位置编码是否能改善Transformer在结构化领域（如生物学、蛋白质或DNA分析）的表现。"

更大更高效的思考

研究人员随后调查了如果PaTH注意力机制更类似地模仿人类认知，即在做决策时忽略旧的或不太相关的信息，会如何表现。为此，他们将PaTH注意力与另一种称为遗忘Transformer（FoX）的位置编码方案结合，该方案允许模型选择性地"遗忘"。由此产生的PaTH-FoX系统增加了以数据相关方式降低信息权重的方法，在推理、长上下文理解和语言建模基准上取得了出色结果。通过这种方式，PaTH注意力扩展了Transformer架构的表达能力。

Kim说，像这样的研究是开发人工智能"下一个重大突破"的更广泛努力的一部分。他解释说，深度学习和生成式AI革命的主要驱动因素是创建"可应用于广泛领域的通用构建块"，例如"卷积层、循环神经网络层"，以及最近的Transformer。展望未来，Kim指出，准确性、表达能力、灵活性和硬件可扩展性等考虑因素一直是并将继续是至关重要的。正如他所说，"现代架构研究的核心事业是试图提出这些新的原语，在保持或改善表达能力的同时，也具有可扩展性。"

这项工作部分得到了MIT-IBM Watson人工智能实验室和Schmidt Sciences的AI2050项目的支持。

Q&A

Q1：PaTH注意力机制与传统RoPE编码有什么区别？

A：PaTH注意力机制使位置信息变得自适应和上下文感知，而不像RoPE那样静态。RoPE为每个词分配基于相对距离的固定旋转，而PaTH注意力将中间词汇视为由数据相关转换组成的路径，能够建模词汇之间路径上含义如何变化。

Q2：PaTH注意力在哪些任务上表现更好？

A：PaTH注意力在多种任务上表现出色，包括跟踪最新"写"命令、多步回忆测试、推理基准、长上下文理解和语言建模。在诊断任务和真实世界语言建模任务中，它都能超越现有注意力机制，同时保持效率。

Q3：PaTH-FoX系统是如何工作的？

A：PaTH-FoX系统将PaTH注意力与遗忘Transformer结合，模仿人类认知中忽略旧的或不太相关信息的能力。它增加了以数据相关方式降低信息权重的方法，在推理、长上下文理解和语言建模基准上取得出色结果。