🤖 machine learning

Interdomain Attention: Beyond Token-Level Key-Value Memory

本文提出了一种名为域间注意力（Interdomain Attention）的新架构，该架构通过核方法将状态空间模型（SSMs）集成到注意力机制中，从而在固定大小的状态上实现查询条件化注意力，进而结合了状态空间模型的可扩展性与长度鲁棒性，以及基于内容匹配的性能优势。

原作者： Naoki Kiyohara, Harrison Bo Hua Zhu, Riccardo El Hassanin, Zhuo Sun, Wenlong Chen, Samir Bhatt, Yingzhen Li

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Naoki Kiyohara, Harrison Bo Hua Zhu, Riccardo El Hassanin, Zhuo Sun, Wenlong Chen, Samir Bhatt, Yingzhen Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用通俗易懂的语言和生动的类比对论文《Interdomain Attention》（域间注意力）的解释。

核心难题：“信息过载”困境

想象你正在尝试写一个故事，但你的记忆力很差。

旧方法（标准 Transformer）： 为了写出下一句话，你必须回顾你写过的每一个字。如果你要写一本 100 页的书，你的大脑就必须同时记住这 100 页的所有内容，以便找到正确的联系。这极其缓慢，且需要巨大的精神能量（计算能力）。随着故事变长，你的大脑会不堪重负。
替代方法（状态空间模型/SSM）： 为了节省精力，你将整个故事总结成一张极小的便签。你只查看这张便签来写下一句话。无论故事多长，这都超级快速且高效。然而，因为你只有一张小小的便签，你往往会遗漏具体的细节。你可能会忘记三章前提到的某个角色的名字。

目标： 作者希望构建一个系统，既拥有便签的速度和效率，又具备通读整本书的记忆力和细节捕捉能力。

解决方案：“Interdomain Attention"（域间注意力）

作者创造了一种名为Interdomain Attention的新方法。你可以把它想象成一位以非常特定方式管理图书馆的聪明图书管理员。

1. “压缩图书馆”（SSM 核心）

图书管理员不是把每一本书（token）都放在书架上，而是使用一台特殊机器（SSM）将整个故事历史压缩成一组固定大小的“摘要系数”。

类比： 想象你有一本 1000 页的小说。与其保留全部 1000 页，不如将故事提炼为 64 个具体的“主题”或“氛围”（例如“英雄的旅程”、“反派的动机”、“场景的基调”）。这 64 个主题就是你的“状态”。无论故事是 10 页还是 10,000 页，你手中永远只需要持有这 64 个主题。

2. “智能查询”（注意力部分）

当你想要写下一句话时，你不仅仅是盲目地查看那 64 个主题。你会提出一个具体的问题（即“查询”）。

魔法技巧： 系统将你的问题翻译成与那 64 个主题相同的“语言”。然后它会检查：“这 64 个主题中，哪一个与我当前的问题最相关？”
结果： 你得到了两全其美的效果。你正在查看一个压缩后的摘要（快！），但你是根据当前的思考内容，从该摘要中挑选出正确的部分（聪明！）。

它是如何工作的（“厨房”类比）

想象你是一位正在做汤（输出）的厨师。

标准注意力： 你有一大锅食材（整个历史）。为了加盐，你必须尝遍锅里每一种食材，以决定加多少盐。随着锅变大，这需要花费永恒的时间。
标准 SSM： 你有一个只有 64 个罐子的小调料架。你直接拿起一个罐子加进去。这很快，但你再也尝不到具体的食材味道了。
Interdomain Attention：
1. 你有一台机器，它根据你到目前为止烹饪的所有内容，不断更新一个固定大小的调料架（那 64 个主题）。
2. 当你想要加盐时，你不需要尝整锅汤。相反，你拿着一把特殊的尝味勺（查询特征图），它能立刻告诉你：“基于当前的味道，你需要混合 30% 的‘番茄’罐和 70% 的‘香草’罐。”
3. 你从固定的调料架中混合这些特定的分量。因为架子很小，所以很快；但因为勺子知道如何根据当前时刻精确混合，所以非常准确。

论文的实际发现

作者在从小型（1.25 亿参数）到大型（13 亿参数）的语言模型上测试了这种新的"Interdomain Attention"。以下是他们的主要发现：

它击败了“小便签”（SSM）： 在每一项测试中，Interdomain Attention 在文本生成方面都优于标准 SSM 方法。它在保持快速的同时，更好地理解了上下文。
在大规模下它击败了“大锅”（标准注意力）： 在他们测试的最大规模（13 亿参数）下，Interdomain Attention 生成的文本实际上比那种查看每个单词的标准方法更好（更低的“困惑度”和更好的常识得分）。
它不会遗忘长故事： 最大的胜利在于长度。当故事长度超过训练范围时，标准方法会变得困惑并犯错。而 Interdomain Attention 即使在故事长度达到训练长度的3.5 倍时，依然保持冷静和一致。它没有变慢或变得更困惑，只是持续工作。
秘密武器： 他们进行了“机制分解”（一种拆解机器以查看哪部分在工作的复杂方法）。他们发现，查询条件投影（将你的问题翻译成摘要语言的“智能勺子”）是成功的主要原因。如果没有这一步，系统就表现得像标准的 SSM，性能较差。

权衡（它的局限性）

论文诚实地指出了其中一个弱点：精确回忆。

类比： 如果你问系统：“第 4 页提到的角色的确切电话号码是多少？”标准方法（查看整本书）非常擅长找到它。而 Interdomain Attention 由于依赖压缩摘要，如果不属于主要“主题”，它在检索确切文本字符串（如电话号码或特定名称）方面表现不佳。
然而，作者指出，这是任何压缩信息的系统（而不仅仅是他们的新方法）的局限性。

总结

Interdomain Attention 是一种 AI 记忆事物的新方法。它不再试图将整个世界都塞进大脑（慢），也不仅仅依靠一张小便签（易忘），而是保持一个智能的、压缩的世界摘要。当它需要写作时，它会使用一个特殊的翻译器来挑选该摘要中完全正确的部分。这使得它快速、高效，并且在处理非常长的故事时，令人惊讶地不易混淆。

技术摘要：跨域注意力

问题陈述

Transformer 和深度状态空间模型（SSM）代表了序列建模中两种对立的架构范式。标准注意力机制（Vaswani 等人，2017）利用不断增长的键值（KV）缓存，通过基于内容的匹配来路由查询。虽然这提供了精确的回忆能力，但其计算成本随序列长度呈二次方增长（ $O(N_q N)$ ），内存消耗呈线性增长（ $O(N)$ ）。相反，深度 SSM（如 S4、Mamba）将整个上下文压缩为固定大小的循环状态，每一步的更新仅需 $O(1)$ 时间。然而，这种压缩的代价是失去了直接的查询 - 键匹配；固定状态无法被特定查询直接寻址，通常导致在需要精确检索或基于内容注意力的任务中表现较弱。

混合架构曾尝试交替使用这两种方法，线性注意力变体则试图直接降低注意力成本。本文提出了一条不同的路径：将 SSM 直接集成到注意力机制中，在保持固定大小状态效率的同时，使每个查询能够关注压缩后的历史。

方法：跨域注意力

核心贡献是跨域注意力（Interdomain Attention），这是一种统一了核注意力与状态空间模型的 Token 混合器。该方法源于对注意力的核回归视角以及 HiPPO 风格的基础投影。

理论推导

核特征映射：注意力核 $K(q, k)$ 通过有限特征映射 $\xi(q)^\top \xi(k)$ 进行近似。
基础投影：该方法不存储所有键和值，而是将键特征 $\xi(k)$ 和值 $v$ 投影到由单个 SSM 递归（具体为 HiPPO 风格）维护的一组共享基函数上。这产生了一个与序列长度无关的固定大小状态。
压缩系数：这些投影的积分生成了三个状态分量：
- $U$ ：键的基础投影系数（诱导变量）。
- $\Gamma$ ：值投影。
- $\eta$ ：归一化常数。
查询条件读取：在每个位置，查询 $q$ 被映射到特征空间 $\xi(q)$ 。输出通过关注压缩系数而非单个 Token 来计算：
$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$
矩阵形式为： $\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$ ，其中 $F_q$ 是查询特征矩阵。

可扩展实现

实际实现是理论推导的学习化松弛版本，旨在实现稳定性和效率：

特征映射：使用 DeltaNet 风格的神经特征映射 $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$ ，其中 SC 是因果深度卷积。
递归核心：使用复对角 S4D 递归来近似 HiPPO 投影。SSM 输入将键特征 $\xi(k)$ 和值 $v$ 拼接。
归一化：在进入 SSM 之前，对键和值部分应用带有可学习每头偏置的输入 RMSNorm。
无分母读取：为避免因输入重缩放导致的不兼容尺度混合，去掉了归一化项 $\eta$ ，从而得到线性注意力读取： $\tilde{O} = F_q U^\top \Gamma$ 。
架构：嵌入到具有预归一化、SwiGLU 前馈和旋转位置嵌入（RoPE）的 Llama 风格解码器中。

主要贡献

构建：一种源自核回归和 HiPPO 基础投影的查询条件固定状态 Token 混合器，清晰区分了理想的理论推导与可扩展的实现。
机制分解：在 1.25 亿参数规模下的消融研究表明，查询条件投影是驱动性能超越标准 S4D 混合器的主要因素。双键/值输入贡献了较小且次要的收益。
扩展研究：在匹配循环状态预算下，对 FineWeb-Edu 数据集进行了从 1.25 亿到 13 亿参数的全面评估。

实验结果

本研究将跨域注意力与仅 S4D Token 混合器及标准 softmax 注意力基线进行了比较。

与 S4D 的性能对比：跨域注意力在所有规模（1.25 亿–13 亿）上均优于 S4D Token 混合器，验证困惑度降低了 13–16%。
与 Softmax 的性能对比：在 13 亿参数规模下，跨域注意力在验证困惑度（7.98 对比 8.63）以及八项常识任务套件（平均准确率提高 3.03 个百分点）上均超越了同配方 softmax 基线。在 1.25 亿参数规模下，其表现与 softmax 大致持平。
长度外推：凭借其固定状态核心，跨域注意力表现出“长度平坦”的行为。虽然 softmax 基线的困惑度在训练上下文（4K）之外显著退化，但跨域注意力在高达 14K（训练上下文的 3.5 倍）的范围内保持了稳定的性能。
回忆局限性：论文指出，与其他固定状态模型一样，跨域注意力在处理精确字符串关联回忆（如“大海捞针”任务）方面，相比 softmax 注意力存在困难；这被识别为固定状态压缩的结构性限制，而非跨域机制特有的缺陷。

意义与主张

本文主张，跨域注意力成功弥合了固定状态 SSM 的效率与基于内容的注意力检索能力之间的差距。通过将特征投影到 SSM 基上，它允许查询关注压缩后的历史，而无需承担标准注意力的二次方成本或线性内存增长。

其主要意义在于证明，将查询条件投影到固定循环状态上，不仅能超越纯 SSM，在更大规模下还能超越标准 softmax 注意力，同时继承 SSM 理想的长度外推特性。作者认为，这种方法为长上下文语言建模提供了一条可行路径，在此类场景中，内存效率以及在变化上下文长度下的一致性表现至关重要。文中提到的未来工作方向包括集成更强的固定状态核心（如 Mamba-3），或将读取机制与快速权重更新规则相结合。