Key-Value Means

本文介绍了键值均值（KVM），这是一种新颖的块循环注意力机制，它通过支持高效且可分块并行训练、具备灵活的状态扩展能力以及次二次方预填充时间，同时仅采用标准运算和极少的额外参数，从而融合了 Transformer 与线性循环神经网络的各自优势。

要点

想象一下，你正在尝试阅读一本非常长的书，但你的大脑只有有限的“工作记忆”来在阅读时把故事保持在脑海中。

当前 AI 的问题
当前的 AI 模型（Transformer）就像一个试图记住书中读过的每一个字的学生。

• 优点：它们极其准确，因为它们面前有整个故事。
• 缺点：随着书变长，它们的“工作记忆”会变得巨大无比。阅读一本 100 页的书只需极少的精力，但阅读一本 1000 页的书则需要耗费大量的时间和能量。这就像试图背着一个每走一步就变得更重的背包。

循环（RNN 风格）模型的问题
RNN 风格的模型采取了一种不同的方法：它们只保留一个小的、固定大小的摘要，并在阅读过程中不断更新它。

• 优点：它们超级快且轻便。无论书有多长，它们的背包永远不会变重。
• 缺点：它们会忘记故事的开头。如果你问它们关于第 10 页的情节要点，它们可能记不住，因为它们只保留着最后几页的内容。

新解决方案：键值均值（KVM）
本文的作者介绍了一种名为**键值均值（Key-Value Means, KVM）**的新方法。将 KVM 想象成一个结合了两者优点的聪明、神奇的笔记本。

以下是它如何通过一个简单的类比来运作：

1. “滑动窗口”（即时上下文）

想象你在读书，你有一个放大镜，只能让你清晰地看到最后几页。这就是“滑动窗口”。KVM 像标准 AI 一样，对最近的词语给予完美的关注。这确保了它不会错过即时的上下文。

2. “压缩摘要”（长期记忆）

当你翻过这几页，旧页面会从你的放大镜中滑出。KVM 不会像 RNN 风格模型那样把它们扔掉，也不会像当前 AI 那样试图携带整本书，而是做了一件巧妙的事：

• 它查看刚刚滑出的那些页面。
• 它问道：“这些页面中哪些是最重要或最独特的？”
• 它将这些重要页面的简短、压缩的摘要写进一个特殊的笔记本中。
• 如果新页面与笔记本中已有的内容非常相似，它就更新现有的笔记；如果它是全新且令人惊讶的内容，它就在笔记本中添加新的一行。

3. “智能合并”（魔法技巧）

本文描述了一种特定的信息合并方式，称为“赢家通吃”规则。

• 想象你有一桶水（新信息）和一块海绵（笔记本）。
• KVM 不是直接把水倒进去，而是找到海绵中与这桶水最匹配的确切位置，并在那里吸收它。
• 它还使用了一种“即时”归一化。想象一下，当你往笔记本里写东西时，你只是简单地记录原始总和和计数，而不去计算平均值。只有当你真正需要读取笔记本里的信息时，你才在最后那一刻进行除法运算（将总和除以计数）来得出正确的平均值。这种“按需”进行的归一化方式，避免了每次写入新条目时都要反复重新计算，从而提高了效率。

为什么这很重要

• 灵活的大小：你可以告诉 KVM 保持一个微小的笔记本（固定大小）以追求速度，或者让笔记本随着书的变长而增长（可扩展大小）。
• 速度与记忆：它允许你选择一个中间地带。你不必在“超级快但健忘”和“超级聪明但缓慢”之间做选择。你可以将其调整为足够快以用于实时场景，同时足够聪明以记住整个故事。
• 无需定制硬件：与某些需要特殊、昂贵计算机芯片才能运行的新方法不同，KVM 可以使用标准计算机和常规软件操作来运行。

结果

作者在语言模型（阅读和书写文本的 AI）上测试了这种方法。

• 短上下文：KVM 的表现与最好的标准 AI 模型一样好。
• 长上下文：当输入扩展到数千个 token 时，带有增长型笔记本的 KVM 比“健忘”的 RNN 风格模型记住了更多的细节，并且比“携带整本书”的全注意力 Transformer 快得多。
• “大海捞针”：在 AI 必须从海量文本中找到特定句子的测试中，允许其笔记本增长的 KVM 版本表现非常出色，证明它能够真正回忆起过去深处的信息。

简而言之，KVM 是一种让 AI 阅读长书而不感到疲惫、不忘记开头、且不需要背负无限变重的背包的新方法。 它通过保持对当下的清晰视野，同时维护对过去的智能、压缩的摘要来实现这一点。

技术摘要

技术摘要：键值均值（KVM）

问题陈述

Transformer 在现代硬件上能够高效训练，但在每个输出 token 的内存和时间消耗上随上下文长度呈线性扩展（$O(N)$ 内存，$O(N)$ 解码时间）。相反，现代线性 RNN（LRNN）实现了每个 token 的恒定内存和时间（$O(1)$），但通常难以处理有限的长上下文回忆。现有试图弥合这一差距的架构往往依赖固定大小的状态（限制回忆能力）或复杂的测试时训练与运行时优化器（影响速度）。因此，亟需一种架构，能够在无需自定义内核或复杂的测试时训练超参数调整的情况下，平衡内存效率、速度和长上下文回忆能力。

方法论：键值均值（KVM）

KVM 是一种新颖的块循环注意力机制，它将块滑动窗口注意力（BSWA）与动态可扩展的压缩状态相结合。它在单个 softmax 注意力层内运行，统一了传统 Transformer（可扩展上下文、分块并行性）和线性 RNN 的优势。

核心机制

1. 带压缩状态的块滑动窗口：
   KVM 以分块方式处理输入。它为一个固定大小的 BSWA 窗口维护最近 token，并为较旧的 token 维护一个单独且定期更新的状态。当一块 token 溢出 BSWA 窗口时，它们会被处理以更新状态，而不是被丢弃。

2. 状态压缩与合并：
   溢出 token 使用一种“赢家通吃”的余弦相似度类合并规则被压缩进状态。

   • 相似度度量： 不同于标准 softmax，KVM 使用最大稀疏更新矩阵（受在线向量量化启发），其中每个溢出键被分配给相关性最强的单个状态键。
   • 即时（JIT）重归一化： 为防止状态向量范数因平均正交或相反向量而随时间缩小，KVM 应用 JIT 归一化。状态键在注意力计算前使用 LayerNorm 进行归一化。状态值被归一化为一个固定的“读取半径”（$\rho_i$），该半径在槽位创建时确定，从而在允许方向变化的同时保留值的大小。
   • 合并门： 一个数据依赖的标量门调节状态吸收的流入溢出键/值的量。

3. 状态扩展策略：
   与固定大小的 RNN 不同，KVM 支持可增长的状态。最“令人惊讶”（冗余度最低）的溢出 token 直接附加到状态中，其余的则被合并。这使得内存增长呈次线性（例如 $O(\sqrt{N})$），同时保持早期上下文的回忆能力。

4. 位置编码处理：
   为了在 BSWA 窗口中保持与旋转位置嵌入（RoPE）的兼容性，同时避免在压缩状态中使用 RoPE（因为压缩状态聚合了来自广泛不同位置的 token），KVM 采用部分 RoPE 置零。状态键的旋转子空间被置零，而 BSWA 窗口保留完整的 RoPE。这使得模型可以在同一次注意力传递中，对状态使用未旋转的查询，对窗口使用旋转的查询。

5. 汇点 token（Sink Tokens）：
   保留一组受保护的初始状态行（汇点），以防止关键早期上下文信息的退化，解决汇点 token 具有不同值大小的问题。

主要贡献

本文提出了以下具体贡献：

• 新颖的块循环公式： 一种使用赢家通吃合并规则将溢出 token 压缩进动态重归一化状态的方法，消除了对单独压缩层的需求。
• 状态扩展策略： 一种将新颖的溢出 token 附加到状态的机制，在保持回忆能力的同时实现次线性内存增长。
• 即时（JIT）重归一化： 一种在即时时刻归一化状态键和值的方案，以维持向量范数并防止在平均过程中产生破坏性干扰。
• 部分 RoPE 共享： 一种通过在状态键中置零 RoPE 维度来在压缩和未压缩区域之间共享位置编码的方法，避免了对额外参数或复杂重新合并机制的需求。
• 统一架构： 单个注意力层，在固定状态 RNN 和全注意力 Transformer 之间进行插值，提供内存效率与回忆能力之间的连续权衡。

实验结果

作者在 Prolong 数据集上训练了模型（1.2 亿和 3.5 亿参数），上下文长度为 8k。

• 长上下文性能：
  • 固定状态 KVM（256 个 token）： 在序列位置损失和短上下文基准测试上优于更大状态的 OVQ/SWA 模型。然而，在极端长度下，面对新颖干扰项（NIAH-S2/S3）的“大海捞针”（NIAH）测试中，其表现挣扎，此时状态容量成为瓶颈。
  • 可增长 KVM（幂律/饱和调度）： "KVM sqrt"变体（状态大小 $\propto \sqrt{N}$）在长上下文基准测试（RULER、LongBench、NIAH）上取得了有竞争力的结果，在 8k 训练上下文之外的外推区域中，匹配或超越了非混合的 GPTAlpha 模型。它在需要检索远距离新颖信息的任务上，显著优于固定状态 KVM 和纯 LRNN（RWKV-7）。
• 短上下文性能： KVM 变体在短上下文基准测试（LAMBADA、ARC、HellaSwag 等）上的表现与标准 Transformer 一致，证实了 BSWA 窗口保留了标准注意力能力。
• 消融研究： 移除值长度归一化导致了最显著的性能下降。移除汇点保护和合并门也大幅削弱了长上下文回忆能力。

意义与主张

本文主张 KVM 成功弥合了固定状态 RNN 与全注意力 Transformer 之间的差距。

• 效率与回忆： 它提供了灵活的状态大小选择，允许用户调整内存效率与回忆能力之间的权衡。使用固定状态时，它提供 $O(N)$ 的块循环行为；使用可增长状态时，它实现次线性内存增长并具备强大的长上下文检索能力。
• 实现简洁性： KVM 可使用标准操作实现，无需自定义内核，并支持分块并行训练和预填充。
• 混合潜力： 该架构可与 LRNN 层结合用于混合解决方案，通过改进的次线性内存增长和长上下文解码能力对其进行补充。
• 无运行时优化器： 与测试时训练（TTT）方法不同，KVM 依赖简单的状态更新规则，而非 SGD 或 Adam 等运行时优化器，从而避免了相关的超参数挑战。

作者得出结论，KVM 证明了以简单有效的方式在固定状态 RNN 和全注意力之间进行平滑插值是可行的，为长上下文建模提供了一个统一的方案。