The Query Knows What to Forget: A Second Erase Direction for Linear Attention

想象一下，你正试图回忆很久以前听过的一个故事。你的大脑能容纳这些记忆的空间是有限的。如果你试图把太多的故事塞进这个狭小的空间，细节就会开始变得模糊，你可能会把一个故事的结尾和另一个故事的开头搞混。这与现代 AI 模型处理长对话的方式有些类似。它们使用一种被称为“线性注意力”（linear attention）的特殊记忆机制，在不耗尽空间的情况下追踪至今为止所说的一切。它们不是将每一个单词都保存成一个巨大的、不断增长的堆积物，而是将信息压缩成一个固定大小的“状态”（state），就像一个会随着新词到来而自动更新的魔法笔记本。

然而，这里有一个问题。随着对话变得越来越长，笔记本上的旧记忆会开始互相干扰。当 AI 试图寻找某项特定信息（比如 10,000 个单词前提到过的一个名字）时，所有其他单词产生的噪声会让它难以检索到正确的答案。科学家们一直试图通过让 AI 变得更聪明地进行“删除”来解决这个问题。他们构建了能够“擦除”旧的、无关信息的系统，以腾出空间给新信息，但直到目前为止，这些擦除器都显得有些笨拙。它们只能根据正在处理的当前单词来决定删除什么，而不是根据 AI 实际在寻找什么。

这篇论文介绍了一个巧妙的新技巧，叫做 QED（查询导向擦除方向，Query-derived Erase Direction），旨在修复这种笨拙感。研究人员在使用具有 3.4 亿个参数的模型进行实验时发现，通过让 AI 的“搜索问题”（查询，query）来辅助决定擦除什么，模型可以更久地保留重要信息。他们在包含 150 亿个 token 的大规模数据集上训练了模型，并发现这种新方法让 AI 能够比以前更久地记住长对话中的细节。具体来说，在一项名为 S-NIAH-1 的测试中，这种新方法使模型能够可靠阅读的文本量翻了一番，将有效极限从大约 8,000 个单词推升到了 16,000 个单词。该论文表明，这是让 AI 更好地阅读长文档迈出的重要一步，尽管作者指出，该工具的某些特定部分可能还需要更多测试，以确定它们是否严格必要。

问题所在：“单向”擦除器

要理解这一突破，我们先来看看这些 AI 模型目前是如何工作的。想象一下，AI 的记忆是一个繁忙的图书馆，书（信息）正在不断被更换。当一本新书到达时，图书管理员（模型）必须决定扔掉哪本旧书以腾出空间。

在本文之前的模型（如 Gated DeltaNet-2）中，图书管理员只根据新书（“键”，key）来决定丢弃什么。这就像是在说：“这本新书是关于猫的，所以我会扔掉那本关于猫的旧书。”这效果还可以，但它有一个盲点。

问题在于，AI 不仅仅是在看新书，它还有一个试图回答的特定问题（“查询”，query）。也许问题是：“主角是谁？”图书管理员需要找到针对那个特定问题的答案。但旧有的“单向”擦除器只知道根据新书的主题来删除内容，而不知道正在被询问的问题。

作者意识到，“干扰”（即让 AI 感到困惑的噪声）恰好就存在于问题所寻找的位置。如果问题是在寻找一个特定的细节，而图书管理员只根据新书的主题来擦除，那么那些令人困惑的噪声就可能正好留在问题的路径上。这就像是在试图于一堆干草中找一根针，而图书管理员只负责移除看起来颜色像针的干草，却忽略了针其实是藏在了一堆稻草下面。

解决方案：“双向”擦除器

论文提出了一个名为 QED 的解决方案。他们没有只使用一个擦除器，而是增加了一个由问题本身引导的第二个擦除器。

可以这样理解：AI 现在拥有一个“键导向擦除器”（旧的）和一个新的“查询导向擦除器”。

1. 键导向擦除器仍然履行其职责：它观察新信息并清除与之匹配的旧内容。
2. 查询导向擦除器则观察 AI 当前正在询问的内容。如果问题是“谁是反派？”，那么这个擦除器会专门针对并移除那些可能遮蔽该问题答案的旧的、令人困惑的信息。

至关重要的是，这个新擦除器的设计是与旧有的擦除器“正交”（orthogonal）的。用通俗的话说，这意味着它清理的是旧擦除器无法触及的部分。这就像是用吸尘器吸走房间角落里扫帚扫不到的灰尘一样。

研究人员通过在 150 亿个 token 的文本上训练模型来测试这一点。他们发现，当他们让“问题”引导擦除过程时，模型在长文本中寻找特定事实的能力变得更强了。

研究发现（以及未发现的内容）

研究结果非常令人鼓舞，但作者也保持了谨慎，没有过度夸大。以下是数据的实际情况：

• 长上下文超能力： 在一项名为 S-NIAH-1 的测试（检查模型是否能在“大海捞针”式的文本中找到目标）中，新方法显著提高了准确率。例如，在 16,000 个 token（单词）的长度下，标准模型的正确率约为 21.8%，而新的 QED 模型达到了 39.8%。在 32,000 个 token 时，差距更加剧烈，从 9.4% 跳升至 15.8%。
• 翻倍的可用长度： 作者指出，这种改进实际上翻倍了可用的上下文长度。此前，模型在 8,000 个 token 左右就开始产生困惑；有了 QED，它在 16,000 个 token 时仍能保持可靠。
• 并非“魔力困惑度”： 有趣的是，这种新方法并没有让 AI 在写作或预测句子下一个词方面变得更好（这是一个被称为“验证损失”的指标）。损失值保持不变。这表明 QED 并没有在一般意义上让 AI 变得更“聪明”；它只是让 AI 在其压缩状态中检索特定记忆的能力变得更强了。

“消融实验”之谜：哪些部分起作用？

为了弄清楚 QED 究竟为什么有效，研究人员进行了实验，通过关闭新系统中的不同部分来进行测试。这就像拆解汽车引擎以观察到底是哪个螺栓让车跑得更快。

他们测试了三个主要组件：

1. 方向（Direction）： 使用“问题”（查询）来引导擦除。
2. 门控（Gate）： 一个决定何时进行擦除的开关（以免删除 AI 仍需要的物体）。
3. 投影（Projection）： 一种数学规则，用以确保新擦除器不会意外破坏旧擦除器的工作。

转折点来了：结果有些复杂。

• 方向是关键： 每当他们使用“问题”来引导擦除时，模型都会变得更好。这部分似乎才是真正的英雄。
• 门控与投影尚不明确： 当他们关闭“门控”或“投影”时，结果并不一致。有时模型表现更好，有时则变差，这取决于他们用于训练的随机种子（起始点）。作者认为，虽然这些部分可能有帮助，但实验并未证明它们是严格必要的。对于他们使用的特定规模的模型，可能是因为“投影”并没有实际发挥作用，因为模型自然地避开了它原本旨在解决的问题。

总结

论文得出结论，QED 是帮助 AI 模型在长对话中记忆事物的成功方法。通过让“查询”协助决定遗忘什么，模型可以清除通常会阻碍其记忆的噪声。

然而，作者对他们的发现局限性保持了诚实。他们并未证明“门控”或“投影”是必不可少的；他们只证明了使用“查询”来引导擦除是最一致的改进方式。他们还注意到，根据具体任务的不同，这种新擦除器的“强度”可能需要进行微调，因为在测试时改变擦除器的“强度”会对不同的模型产生不同的结果。

简而言之，论文表明，如果你想让 AI 阅读整部小说并在最后一章记住反派的名字，你应该给它一个“查询引导的擦除器”。但你可能仍需进行更多测试，以弄清楚如何调节这台新机器上的其他旋钮。