主题 33：Slim Attention、KArAt、XAttention 和 Multi-Token Attention 详解 – Transformer 模型中真正改变了什么？

社区文章发布于 2025 年 4 月 4 日

🔳 我们探索了四种先进的注意力机制，它们改进了模型处理长序列的方式，减少了内存使用，并使注意力可学习。

人工智能中的注意力是一种基本技术，随着我们继续使用 Transformer 等架构，它将永远是一个热门话题。注意力机制让我们得以一窥模型在做决策时关注的重点。它们允许模型动态地关注其输入的特定部分，研究人员正试图利用注意力权重进行可解释性研究，并找出模型做出选择的原因。

两种作为 Transformer 核心机制的注意力类型，曾彻底改变了人工智能模型的效率：1) 自注意力 (Self-Attention)，它允许每个 token “查看”序列中的所有其他 token 以理解上下文；2) 多头注意力 (Multi-Head Attention, MHA)，它并行运行多个注意力机制以捕获不同类型的关系。现在，它们已成为所有主要当前大型语言模型（如 GPT、BERT、T5 和 LLaMA）的基础。

另一个值得注意的注意力机制是 DeepSeek 的 多头潜在注意力 (Multi-Head Latent Attention, MLA)，我们曾在之前的一期中介绍过。它超越了 MHA，通过修改 MHA 机制将 KV 缓存压缩成更小的形式，从而减少内存使用。

即使是这些例子也表明，新的注意力技术 = 新的可能性和能力。它们也为引导或指导生成开辟了可能性。最近，我们观察到研究人员越来越关注人工智能中的注意力。这给了我们一个提示，即社区正在寻求新的机制，将我们日常使用的模型提升到一个新的水平。今天，我们将深入探讨四种最新的多样化注意力机制：1) Slim Attention，它可以更快地处理长上下文，并将内存使用减少多达 32 倍；2) XAttention，提高稀疏注意力在长序列（包括文本和视频）中的有效性；3) Kolmogorov-Arnold Attention (KArAt 和 Fourier-KArAt)，一种完全不同的方法，专注于使注意力可学习和适应性强；以及 4) Multi-Token Attention，它允许词语“组队”来决定什么才是重要的。它们如何工作？哪些模型可以从中受益？我们是否引起了你的注意？很好。让我们开始吧。你将学到很多！

📨 点击关注！如果你想直接在收件箱中接收我们的文章，请在此订阅

什么是 Slim Attention？

处理长上下文对于所有大型语言模型来说仍然是一个严峻的挑战，这也会占用大量内存并减慢所有操作——尤其是在推理过程中生成新 token 时。OpenMachine 的研究人员决定克服这个问题，他们专注于注意力机制，并提出了 Slim Attention。这项技术可以在使用更少内存的情况下更快地获得与使用 MHA 相同的效果，这对于扩展大型模型非常有用。例如，在 Whisper 等模型中，Slim Attention 可以将内存使用减少 8 倍，并在运行大型批次时将文本生成速度提高多达 5 倍。在某些情况下，例如 T5-11B 模型，它甚至可以将内存使用减少 32 倍！

让我们更精确地看看 Slim Attention 如何实现这些令人印象深刻的结果 →

Slim Attention 如何工作？

关键在于使用与 MHA 相同的数学原理，但内存使用更高效。当模型处理长文本块时，它们需要记住很多东西——特别是 K（键）和 V（值）——这在一个叫做 KV 缓存的地方，它可能会变得非常大。

通常，当 Transformer 模型处理文本时，它们会在内存中存储键 (K) 和值 (V)。但是 Slim Attention 发现了一个巧妙的技巧。

由于 K 和 V 来自相同的输入并使用方阵，它可以反向计算并从 K 重建 V，而不是同时存储两者，这会产生完全相同的结果。从 K 重建 V 有两种选项：

重新计算 V 从 K，然后进行注意力计算。这很简单，但需要更多的计算。
先进行注意力计算，然后应用 V-from-K 变换。这更高效且快速，尤其是在一次生成一个 token 时，但实现起来稍微复杂。

选项 2 通常在生成期间更好，因为它在节省计算的同时仍能减少内存使用。

图片来源：Slim Attention 原论文

因此，Slim Attention 不再存储 K 和 V，而是只存储 K，并在需要时重新计算 V。这使得内存使用量减半，或者让模型能够一次“看到”的文本量增加一倍。

研究人员还解决了应用 RoPE（旋转位置嵌入）的潜在问题，RoPE 有助于模型理解每个 token 在序列中的位置。问题是 RoPE 只应用于 Q（查询）和 K（键）——而不应用于 V（值）。因此，要用 Slim Attention 从 K 重建 V，你可能需要先撤销 RoPE。根据是否使用稀疏性技巧，有两种方法可以做到这一点：

如果我们不使用稀疏性，我们会在缓存中存储原始 K，并在从内存中读取 K 后再应用 RoPE。这里，原始 K 用于计算 V。
如果我们使用稀疏性，我们存储已进行 RoPE 处理的 K，但对于实际使用的少数 K 向量，我们会撤销 RoPE。这称为 RoPE 解码，幸运的是，它很简单，并且使用与编码相同的正弦/余弦数学。

Slim Attention 凭借其巧妙简单的 KV 处理方法取得了哪些成就？

Slim Attention 的优势

它能让模型更快，通常能将内存需求减少一半，而不损害准确性——这对于长输入和大批量尺寸尤其有用。
- 对于单个输入（批量大小为 1），它可以将速度提高多达 1.8 倍。
- 对于多个输入（批量大小为 16），速度提升可达 2 倍。
- 而且，正如我们已经提到的，在 Whisper 等模型中，它可以将内存使用减少 8 倍，速度提高多达 5 倍。对于 T5-11B，它可以将内存使用减少 32 倍。
运行它更便宜。
Slim Attention 易于添加到使用多头注意力（MHA）的现有模型中，例如 CodeLlama、Aya、Phi-3-mini、SmolLM2、像 LLAVA 这样的视觉-语言模型、像 Qwen2-Audio 这样的音频-文本模型，以及像 Whisper 和 T5 这样的编码器-解码器模型。无需重新训练。
如果模型中使用偏置，Slim Attention 可以轻松支持偏置。

尽管 Slim Attention 具有这些令人印象深刻的优势，但它也有一些可能不明显的缺点。

并非没有局限性

虽然 Slim Attention 节省了内存，但它在运行时重建 V 时会多消耗一点计算资源。在计算是瓶颈而不是内存的系统上，这可能会减慢速度而不是提供帮助。
它不适用于一些较新的注意力变体，如 MQA（多查询注意力）或 GQA（分组查询注意力），这些变体被许多新的大型语言模型用于提高效率。
RoPE 处理增加了复杂性，尤其是在使用稀疏性技巧时。
对于小型模型和输入，它带来的收益较小，因为内存使用本来就很低。
Slim Attention 不适用于非方阵权重矩阵——当模型（如 T5-11B）的投影权重中输出大小 > 输入大小时。反转这些大矩阵更困难或更昂贵。
不适用于所有位置编码方法。

总之，在内存使用是主要瓶颈且我们需要更快地处理 token 的情况下，Slim Attention 是 MHA 的一种透视替代方案。

处理长上下文和多样化数据（如视频、图像和音频）是当今大型语言模型面临的关键挑战。新的注意力机制正在不断涌现以解决这一问题。我们接下来要探讨的机制采用了不同的方法。

什么是 XAttention？

XAttention 的主要思想

加速 AI 模型处理任何类型长输入的一种方法称为块稀疏注意力。它只关注输入中最重要的部分，而不是检查所有可能的连接。但找出哪些部分是重要的通常也需要大量时间和精力，这会抵消速度优势。

这就是 XAttention 出现的原因，它使块稀疏注意力效率更高。它由清华大学、麻省理工学院、上海交通大学和 NVIDIA 的研究人员推出。它不进行繁重的计算来查找重要部分，而是“查看”一些简单的东西：注意力矩阵中对角线（从左下到右上）上的值之和。这被证明是发现最重要内容的绝佳捷径。

XAttention 通过快速估计哪些部分重要，运行速度更快，在某些情况下甚至能将速度提高 13.5 倍，同时保持与常规注意力相同的准确性。让我们一步一步地看看它是如何精确工作的。

XAttention 如何工作？

XAttention 首先预测重要性。其核心思想是关注反对角线——每个块中从左下到右上的线条。通过对这些线条上的值求和，模型可以快速估计块的重要性。

图片来源：XAttention 原论文

以上图为例：输入被分成 8×8 的块。红线表示求和值较高（更重要），蓝线表示求和值较低（不那么重要）。但这为什么有效呢？因为每个 token 至少贡献给一条反对角线。此外，高影响力的模式——比如垂直线或斜线——倾向于与这些线相交。这为 XAttention 提供了可靠、快速的优先级信号。

图片来源：XAttention 原论文

选择要保留的块 一旦使用反对角线对块进行评分

它们的得分通过 softmax 进行归一化，将其转换为概率。
根据这些分数，模型“选择”只保留得分较高的块，并对这些选定的块执行注意力操作（在上面的第一张图像中用红色标记）。它只是最小的一组块，其组合重要性得分超过某个阈值。

为每个注意力头调整阈值 模型的不同部分，称为注意力头，表现不同——有些需要更多细节，有些可以更稀疏。XAttention 使用动态规划策略来为每个头单独调整阈值。它在每一步缓慢地将阈值降低 10%，以找到速度和准确性之间的正确平衡。这部分有助于进一步微调性能，但它是可选的。XAttention 在没有它的情况下仍然表现出色。

现在分析这种注意力机制能够让模型实现什么，以及它的主要优势，这很重要。

XAttention 性能的优势

XAttention 在与其他稀疏注意力机制相比，展现出令人印象深刻的性能提升，在某些情况下，其得分与完全注意力相当。

效率和速度：它在长序列上提供高达 13.5 倍的注意力计算速度，并且由于模式更轻，它也优于其他更短输入的处理方法。这大大降低了内存和计算成本。
轻量级设计： XAttention 的模式选择速度提高多达 24.9 倍。
改进的文本理解：它在长文本任务上优于其他稀疏注意力方法，即使在非常长的序列长度下也能保持准确。
- 在 RULER 上，它甚至可以击败完全注意力。
- 在 LongBench 上，它在所有任务中的平均准确率最高。

图片来源：XAttention 原论文

顶尖的视频理解： XAttention 可以处理长达 1 小时的视频，每秒 1 帧，击败了其他稀疏注意力方法甚至 FlashAttention。
视频生成效率：它非常接近完全注意力，计算节省超过 50%。
即插即用： XAttention 可以轻松添加到现有 Transformer 模型中，无需进行重大更改。
可扩展性： 专为长上下文 Transformer（LCTM）设计，它帮助模型扩展到 256k+ token 或小时长的视频。

总的来说，XAttention 使在现实世界中更容易使用强大的模型，而无需在计算上投入巨资。其轻量级设计和即插即用特性使其在现实世界的 AI 系统中具有实用性——特别是那些涉及多模态 AI 的系统。然而，它仍然存在一些我们应该考虑的问题。

需要注意的局限性

反对角线评分是一种巧妙的捷径，但它仍然是一种近似。在某些罕见情况下，它可能会遗漏重要的模式。
稀疏性与准确性的权衡：选择过高的阈值或过大的步长可能会使注意力过于稀疏，从而导致准确性下降。
视频生成需要预热：在视频生成过程早期应用 XAttention 会导致布局问题。这可以通过在最初的 5 步中使用完全注意力添加“预热阶段”来解决。然而，这增加了集成复杂性。
专为块结构设计：XAttention 结构可能不适用于不规则或高度动态的输入。
虽然它对许多模型来说是即插即用的，但尚未直接测试所有 Transformer 变体。使用非标准注意力格式的模型可能需要调整。

尽管存在这些问题，XAttention 被证明是稀疏注意力机制的强大升级，在长序列上表现出有益的性能。

现在，是时候转向第三种注意力类型——Kolmogorov-Arnold Attention，并探索它如何改进注意力机制。

什么是 Kolmogorov-Arnold Attention (KArAt)？

在注意力机制中应用 KAN 的思想

中佛罗里达大学的研究人员采取了一种完全不同的方法，这是一种更具哲学性的方法——他们旨在探索 Kolmogorov-Arnold Networks (KANs) 如何在视觉任务中增强 Transformer，不是通过提高效率，而是通过在学习方式上更智能、更灵活。他们主要的问题是：如果注意力本身可以学习呢？

KANs 快速回顾： KANs 不会用权重乘以输入，而是对每个输入维度应用学习到的样条函数。然后，就像 MLP 一样，它们将这些转换后的值求和并传递到下一层。这有助于 KANs 理解数据中更复杂的模式。训练后，研究人员可以查看模型内部，了解它学到了什么，这几乎就像窥视其“思考”过程一样。

KANs 在发现数学公式和随时间学习简单的一维模式等任务中表现出强大的结果。然而，将 KANs 与 Transformer 等更强大的架构结合的测试并不多。

为了探讨基于 KANs 的可学习注意力机制是否可以取代视觉 Transformer (ViTs) 中传统的不可适应的 softmax 并做得更好，研究人员创建了一个版本，称为 Kolmogorov-Arnold Attention (KArAt)，后来又创建了一个更实用、模块化的版本，使用傅里叶特征，称为 Fourier-KArAt。让我们看看这些机制是如何工作的。

KArAt 如何工作？

在视觉 Transformer (ViT) 中，每个注意力头“查看”输入并“决定”关注哪些部分，从而构建一个注意力矩阵，其中每一行表示输入的一部分对所有其他部分的注意力程度。每行都通过 softmax 函数运行，将原始数字行转换为概率（所有数字都变为正数，并且它们之和为 1）。

KArAt 用使用 KANs 的自定义可学习激活函数取代了 softmax。结果是，模型学习了它自己的数据转换。

这种方法受到 Kolmogorov-Arnold 定理的启发，该定理解释了复杂函数如何分解为更简单的函数。在这里，KArAt 使用基函数构建函数，基函数就像数学构建块，可以是：

傅里叶基：正弦和余弦波的组合。
B 样条：平滑的曲线拟合函数。
小波或分形用于更专业的用例。

图片来源：Kolmogorov-Arnold Attention 原论文

每个注意力行都通过这些自定义函数之一，这些函数由混合不同的基函数组成。模型学习每个基函数的使用量——这些就是它在训练期间调整的参数。

但是输出可能不会像概率一样。因此，特殊的 ℓ1 投影通过使值变为正并使它们总和为 1 来解决这个问题。它还确保注意力仍然有意义。

但这里出现了现实问题：计算成本极高。解决方案是使实际架构更小、更智能，这就是 Fourier-KArAt。

Fourier-KArAt 如何工作？

注意力矩阵通常具有低秩结构，这就是为什么它们不需要使用所有维度来表示重要信息。因此，研究人员提出缩小可学习算子的大小。

与其使用一个具有 N × N 个值的完整矩阵，我们可以使用一个更小的 r × N 矩阵（其中 r 也远小于 N）。
这个较小的矩阵将注意力数据压缩到较低维空间。
然后，另一个矩阵将其投影回原始大小。这个技巧显著减少了内存使用。

这个新的轻量级版本被称为 Fourier-KArAt（带有傅里叶基的 Kolmogorov-Arnold 注意力）。

为什么傅里叶基函数优于其他选项？很简单。它们可以表示任何平滑的周期函数，对 GPU 友好，并且扩展性更好。

Fourier-KArAt 可以用于实际的 ViT 模型，其可学习的注意力模块可以通过两种方式进行配置：

图片来源：Kolmogorov-Arnold Attention 原论文

分块式： Transformer 中的每一层（块）都有自己独特的注意力函数。这能捕获更详细的模式，但会使用更多参数。
通用式：所有层共享相同的注意力函数，这使模型更小，可能更适合更简单的任务。

然而，这两种设置在灵活性和效率之间进行权衡，具体取决于任务。但是 Fourier-KArAt 有什么性能优势吗？

Fourier-KArAt 的优势

以下是 Fourier-KArAt 显示改进的主要场景：

小型模型上更好的准确性：在 ViT-Tiny 中，Fourier-KArAt 在无需额外技巧的情况下提高了准确性：
- CIFAR-10 上提高 +5.4%
- CIFAR-100 上提高 +7.4%
更清晰的注意力图： 可视化显示，Fourier KArAt 倾向于更精确地聚焦于图像的重要部分，帮助模型做出更清晰的决策。

我们还来总结一下它的总体优势：

可学习注意力： Fourier-KArAt 赋予模型更大的灵活性，以适应其聚焦输入数据的方式。
傅里叶平滑近似： 使用正弦和余弦函数有助于模型平滑地近似类似 softmax 的行为，同时保持可训练和对 GPU 友好。
通过低秩设计减少计算： 它使用注意力机制的低维版本，而不是繁重的全尺寸操作，从而降低内存和计算成本。
对不同架构灵活： 该设计可能适用于任何数学基础，这使得将来可以轻松尝试不同的注意力行为。

然而，Fourier-KArAt 的表现并不如我们所愿。

并非没有局限性

在大型模型中没有明显优势：在 ViT-Small 和 ViT-Base 中，准确性提升微乎其微或不一致，有时甚至比原始模型更差。
Fourier-KArAt 的实现比常规注意力更复杂：解释注意力图更困难，可视化可能需要额外的调整。
它增加了大量额外参数，这会使模型更难训练和调整，尤其是在更深或更大的架构（如 ViT-Base）中。
Fourier-KArAt 可能导致训练路径更不稳定，尤其是在大型模型中。
虽然它在开始时学习速度快，但训练在后期阶段会减慢。
高内存使用：使用 Fourier-KArAt 进行训练会消耗大量 GPU 内存，在某些情况下，单个图像甚至高达 60 GB。

Fourier-KArAt 目前的局限性多于其优点，但它为注意力提供了一种全新且替代的方法，具有巨大的潜力。像许多其他发展一样，Fourier-KArAt 需要更多的时间和实验才能达到更高的质量水平。这将使其能够充分利用其现有优势并获得新优势。

Multi-Token Attention (MTA)

最近，Meta 的 FAIR 研究人员引入了另一种先进的注意力机制——Multi-Token Attention (MTA)，它在长上下文任务中特别有效。MTA 不再只将一个“查询”词与一个“键”词进行比较，而是让模型同时考虑小范围的相邻词组。这种注意力信息还在不同的注意力头之间共享。以下是 MTA 如何实现这一概念。

MTA 如何工作？

MTA 实现了图像处理中的卷积技术，其中附近的像素相互影响。但在 MTA 的情况下，它是 token 而不是像素。MTA 在两个层面使用这种技术：

图片来源：MTA 原论文

键-查询卷积

模型在决定关注哪些内容时，会查看附近的词组。这就像一个滑动窗口，允许模型将词语作为一个整体而不是单独地进行查看。

例如，如果模型必须回答“爱丽丝在哪里看到了兔子？”这个问题，它需要找到一个同时提到“爱丽丝”和“兔子”的句子。通过键-查询卷积，模型可以将“爱丽丝”和“兔子”的 token 组合起来，以找出它们在句子中同时出现的位置，而标准注意力只能找到它们单独出现的位置。

有两种方法可以应用这个过程：Softmax 之前

在模型选择哪些词是重要的之前（softmax 之前）。
在模型选择哪些词是重要的之后（softmax 之后）。

头混合卷积

MTA 允许注意力头共享它们所发现的内容。例如，一个寻找“爱丽丝”的头和另一个寻找“兔子”的头可以协同工作，这样头组就能找出两者同时出现的地方。这种混合可以在 softmax 之前或之后进行。

为了改善学习并避免训练问题，MTA 使用带深度缩放的组归一化。MTA 中的每个注意力头都应用自己的组归一化，而不是在整个层中共享一个。缩放因子根据层的深度而变化，因此，在更深的层中，模型使用更强的归一化来保持稳定。这有助于控制归一化在每一步对学习的影响程度。

MTA 的优势

MTA 更善于结合线索来找到正确答案。
更强的上下文理解：MTA 允许注意力同时依赖多个 token，而不仅仅是一个键-查询对，因此它可以更好地定位复杂或复合信息。它在不需要额外参数的情况下提高了语言理解能力。
结合各头的优势，从而实现更好的推理。
即使存在多个隐藏事实，也能表现良好。
自动学习注意力模式。
具有灵活性，可以针对速度与准确性等不同权衡进行调整。
长上下文准确性提升：MTA 在需要从长文本中精确检索的任务上表现更好。
解决复杂的注意力任务：在标准 Transformer 会失败的玩具任务（例如识别包含多个目标字母的块）中，MTA 成功实现了零错误。
开销小：性能提升带来的参数增加微乎其微（仅约 0.001%）——因此它在模型大小方面是高效的。