🤖 machine learning

Cascade Token Selection for Transformer Attention Acceleration

本文提出了一种级联令牌选择机制，通过在各层继承并增量更新代表性令牌来加速 Transformer 注意力计算，从而在保持高信息保留率的同时，将选择复杂度从 $O(T^2 d)$ 降低至 $O(T r d)$ 。

原作者： Stephen J. Thomas

发布于 2026-05-06

📖 1 分钟阅读☕ 轻松阅读

原作者： Stephen J. Thomas

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

想象一个 Transformer 模型（现代人工智能的大脑）是一座巨大的多层图书馆。每当 AI 阅读一个句子时，它都会派遣一支“图书管理员”团队（层）上楼来整理信息。

在标准图书馆中，书架上的每一本书（token）都需要与其他每一本书进行比较，以寻找关联。如果你有 512 本书，仅为了决定哪些书重要，就需要进行超过 260,000 次比较。这既缓慢又昂贵，尤其是对于长故事而言。

问题：“重新检查”瓶颈

一种名为 ADA 的先前方法试图解决这一问题。它意识到，大多数书实际上只是几本“关键”书的副本或非常相似。ADA 不再比较所有 512 本书，而是挑选一小部分“代表性”书籍（例如 200 本），并忽略其余部分，假设它们是冗余的。

然而，ADA 有一个隐藏成本：为了找出这 200 本关键书籍，它必须在图书馆的每一层，从头开始重新检查每一本书与所有其他书籍的关联。 这就像在每一层楼都雇佣一支新的图书管理员团队来重新整理整座图书馆，尽管从下一层楼上来后，书籍并没有发生太大变化。寻找关键书籍的成本几乎与阅读它们的成本一样高。

解决方案：“级联”电梯

本文介绍了一种巧妙的捷径，称为级联 Token 选择（Cascade Token Selection）。

将图书馆的楼层想象成 AI 中的层。作者发现了一个令人惊讶的事实：第 10 层的“关键书籍”组与第 11 层的组几乎完全相同。 在一层楼上重要的书，在下一层楼依然重要。AI 不会因为向上移动了一层楼，就突然决定某本随机书籍变得重要。

级联方法不再在每一层重新检查整座图书馆，而是这样做：

继承：它从下一层获取“关键书籍”列表。
验证：它仅检查这些特定的关键书籍是否仍然是关键，以及是否有任何被“忽略”的书籍突然变得重要。
更新：它进行微小的调整（添加或移除几本书），而不是从头开始。

类比：音乐会人群

想象一场音乐会，人群就是 AI 的数据。

旧方法（独立选择）：在每一首歌时，一名保安扫描整个由 10,000 人组成的人群，以找出 500 名最兴奋的粉丝。这花费了漫长时间。
新方法（级联）：保安查看上一首歌中 500 名兴奋粉丝的名单。他知道其中大多数人仍然兴奋。他只需检查这 500 人是否仍然兴奋，以及后排是否有任何新人突然跳了起来。他不再重新扫描整个人群。

结果：论文发现了什么

作者在强大的计算机芯片上，针对三种不同的 AI 模型（GPT-2、GPT-J 和 OPT）测试了这种方法。结果如下：

巨大节省：通过不再每次重新扫描整个人群，他们节省了22% 到 63% 的计算机工作量，这些工作量原本仅用于寻找重要 Token。模型越深（楼层越多），节省幅度越大。
稳定性：“关键书籍”列表从一层到下一层保持了 83% 到 94% 的相同性。这证明了随着 AI 深入，其对重要性的理解非常稳定。
安全性：该方法是“保守”的。它绝不会意外丢弃一本真正重要的书。它可能会保留几本额外的“可能”书籍（使列表稍大），但它保证绝不会遗漏任何关键书籍。这意味着 AI 的回答依然同样准确。

为什么这很重要

论文总结认为，之所以有效，是因为 AI 内部的“世界观”随着深入而平滑变化。这不是混乱的跳跃，而是温和的演变。通过利用这种平滑性，级联方法将沉重、缓慢的过程转变为轻盈、快速的过程。

简而言之：不要在每一步都重新发明轮子。只需检查你正在滚动的轮子是否仍然是圆的，如果不是，就修正微小的晃动。 这使得运行大型 AI 模型显著更快、更便宜。

技术摘要：级联令牌选择以加速 Transformer 注意力机制

问题陈述
Transformer 模型在自注意力机制中面临显著的计算瓶颈，其中计算 $QK^\top$ 乘积及随后的 softmax 加权求和的成本随序列长度呈二次方增长（ $O(T^2 d_h)$ ）。尽管 FlashAttention 等近期方法优化了内存带宽，但算法复杂度仍是长序列处理的障碍。

配套论文《激活去相关注意力（ADA）》提出了一种解决方案，通过在每一层选择一小部分 $r \ll T$ 个“代表性”令牌，将复杂度从 $O(T^2)$ 降低至 $O(r^2)$ 。ADA 利用 Gram 阈值识别这些令牌：若某令牌与所有前序令牌的最大余弦相关性低于阈值 $1 - \tau^2$ ，则被视为代表性令牌。然而，选择过程本身需要在每一层计算完整的 $T \times T$ Gram 矩阵，产生 $O(T^2 d)$ 的成本。对于中等长度的序列（例如 $T=2048$ ），该选择开销与其旨在降低的注意力成本相当，从而抵消了潜在的加速效果。本文解决的核心问题是 ADA 中代表性令牌选择步骤的高计算成本。

方法论：级联机制
本文引入了一种利用代表性令牌集深度相干性的“级联”机制。作者观察到，第 $l$ 层的信息令牌集对第 $l+1$ 层的令牌集具有强预测性。这一观察在理论上基于 Transformer 块的 Lipschitz 连续性：在一层中为近重复（冗余）或去相关（代表性）的令牌，在添加上下文信息后，倾向于在下一层保持这种关系，仅产生微小扰动。

级联机制不再在每一层重新计算完整的 $T \times T$ Gram 矩阵，而是：

继承来自上一层的代表性集合 $R^{(l)}$ 。
验证继承的集合，通过计算两个较小的 Gram 乘积：
- 在继承的代表性令牌之间计算 $r \times r$ 乘积，以检查是否有任何令牌变得冗余。
- 在非代表性令牌与有效的继承代表性令牌之间计算 $(T-r) \times r$ 交叉乘积，以检查是否有任何非代表性令牌变得去相关（即成为新的代表性令牌）。
基于这些检查，通过少量的添加和移除操作更新集合。

该方法将每层的选择复杂度从 $O(T^2 d)$ 降低至 $O(Trd) $，加速因子为$ T/r$。

主要贡献

算法创新：提出了一种级联令牌选择算法，用增量验证和更新步骤替代完整的 Gram 矩阵计算。
理论依据：证明了（命题 3）代表性集合的深度相干性是 Transformer 块 Lipschitz 连续性及残差连接加性性质的结果。
保守性保证：证明了（命题 4）级联方法是保守的；它生成独立代表性集合的超集。它绝不会遗漏真正的代表性令牌，从而确保注意力质量不下降，尽管可能会略微增加代表性令牌的数量（ $r$ ）。
正交加速：本文将该方法定位为 门控子空间推理（GSI） 在令牌维度上的类比，后者通过深度相干基继承加速线性层。两者共同覆盖了 Transformer 的前向传播全过程（包括 $d$ 轴和 $T$ 轴）。

实验结果
该方法在三个模型系列（GPT-2 124M、GPT-J 6B、OPT 6.7B）上进行了验证，使用 AMD MI300X 硬件，序列长度 $T=512$ ，Gram 阈值 $\tau=0.30$ 。

深度相干性：深层中连续层代表性集合之间的 Jaccard 指数较高（0.83 至 0.94），证实了令牌集的稳定性。
Gram 运算节省：级联机制减少了总 Gram 运算量：
- GPT-2（12 层）：22%
- GPT-J（28 层）：52%
- OPT（32 层）：63%
  节省量随模型深度增加而增加，因为第 0 层的初始完整 Gram 成本被分摊到了更多层上。
更替率：在深层中，更替率（集合被修改的比例）降至 3–18%，表明继承的集合仅需微调。
扩展预测：作者预测，对于更长的序列（例如 $T=8192$ ），其中 $r$ 呈亚线性增长，选择加速比可达 $18\times$ ，将选择成本从主导项降低为可忽略项。

意义与主张
本文声称，级联机制揭示了训练后 Transformer 的一种结构特性：信息令牌集是一种相干的结构属性，会通过网络深度传播。通过利用这种相干性，该方法使 ADA 的选择开销变得实用，使得注意力压缩能够实现其理论加速，而不再受选择步骤的瓶颈限制。

作者强调，这是对现有方法（如优化内存带宽的 FlashAttention 和使用固定模式的稀疏注意力）的互补加速。与永久改变表示的令牌剪枝或合并方法不同，级联 ADA 根据输入结构动态选择代表性令牌，无需架构更改或重新训练。GSI 与级联 ADA 的结合提供了推理成本的乘积级降低，有效地将前向传播的复杂度从名义模型规模降低为输入的“有效复杂度”。