🤖 machine learning

A Faster Generalized Two-Stage Approximate Top-K

本文通过在每个分区中选择前 $K'$ 个元素而非仅前1个元素，对两阶段近似Top-K算法进行了推广，从而提供了更紧的理论召回率上界，并在保持相同预期召回率的同时，在Cloud TPUv5e上实现了数量级的加速。

原作者： Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

发布于 2026-05-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

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

想象一下，你是一家拥有数百万本书籍（数据）的巨型图书馆的管理员。每天，你都需要找出前 K 本最受欢迎的书籍（即 K 个最大数值），以便向访客推荐。

在计算机芯片（特别是用于训练巨型 AI 模型的芯片）的世界里，找出这些“最受欢迎”的项目竟然既缓慢又昂贵。这就像试图通过一本一本地阅读每一本书来找出前 100 本好书，尽管你的图书馆本设计为能够同时对巨大的书籍堆进行数学运算。

以下是本文为解决该问题所做的简要分解。

旧方法：“逐个”筛选

一种先前的方法（Chern 等人，2022 年）试图通过两步流程来加速这一过程：

分割：想象将你的图书馆分成 100 个不同的房间（桶）。
首次扫描：在每个房间里，助手仅挑选一本最受欢迎的书并将其带到前台。
最终排序：管理员随后仅查看那 100 本书（每个房间一本），并从中选出总体前 100 本。

问题所在：这种方法过于谨慎。由于每个房间只挑选了一本最好的书，它经常遗漏了藏在同一房间内的第二或第三好的书。为了确保不遗漏任何内容，他们不得不使用许多房间（桶），这意味着管理员在最终仍须对一大堆书籍进行排序。这仍然太慢了。

新想法：“前 K"筛选

本文的作者意识到，计算机芯片拥有未被利用的额外算力。他们提出了第一步的更智能版本：

助手不再仅从每个房间挑选第 1 名的书，而是从每个房间挑选**前 K'**本书（例如，前 4 本）。

为什么这更好？

需要的房间更少：因为助手从每个房间抓取了更多的书，你就不需要那么多房间来确保捕捉到所有热门书籍。
排序更少：尽管助手从每个房间抓取了更多的书，但发送给管理员进行最终排序的书籍总数实际上要小得多。
结果：管理员需要排序的是一小堆书，而不是一座山。

硬件的“魔力”

本文解释说，现代计算机芯片（如 Google 的 TPU）就像拥有不同工作站的巨型工厂：

矩阵单元（MXU）：一个超高速工厂，擅长繁重的数学运算（乘法），但不擅长排序。
向量单元（VPU）：一个较小、较慢的工作站，擅长排序和挑选优胜者。

旧方法浪费了 VPU 的时间。新方法利用 VPU 在 MXU 忙于数学运算的同时抓取“前 K'"本书。这就像让一名工人在机器仍在运行时从传送带上抓取最佳物品，从而无需等待时间。

结果：加速 AI

作者在 Google TPU 芯片上对此进行了测试：

旧方法：找出前几名书籍耗时很长，往往比最初生成该列表的数学运算还要慢。
新方法：通过从每个桶中抓取“前 4 名”而不仅仅是“第 1 名”，他们将最终排序的工作量平均减少了7 倍。
融合：他们甚至成功将“挑选”步骤与“数学”步骤结合起来，使它们在同一时刻发生。

核心结论：
在现实世界测试中（在大型 AI 模型中寻找前 2% 的数据），他们的新方法使该过程比之前的标准快了24 倍。这意味着 AI 模型可以更快地进行训练和运行，而不会损失准确性。

总结类比

旧方法：你有 1,000 个团队。每个团队向你派出他们的最佳球员。然后你必须面试 1,000 名球员以找出前 100 名。
新方法：你拥有更少的团队（例如 250 个）。每个团队向你派出他们的前 4 名球员。你只需面试 1,000 名球员（250 个团队 × 4 名球员），但由于你从每个团队获得了更多选择，你找到真正最佳球员的可能性与之前一样，而且由于你更好地组织了团队，你完成得更快。

本文从数学上证明，这种“前 K'"方法不仅仅是一个猜测；它是一种保证能以显著更少的工作量获得相同质量结果的可靠方法。

技术摘要：一种更快的广义两阶段近似 Top-K

问题陈述
在数组中识别前 K 个元素是许多机器学习算法中的基本操作，包括最大内积搜索（MIPS）、K 近邻（KNN）以及大型基础模型的优化（例如用于分类 logits、MLP 块、注意力机制中的稀疏性，以及 KVCache 压缩）。在现代加速器（如 TPU 和 GPU）上，计算精确的 Top-K 往往成为性能瓶颈。这些设备针对稠密矩阵乘法（MXU/TensorCores）进行了优化，而 Top-K 操作通常依赖于标量/向量单元（VPU/CUDA 核心），其速度可能显著慢于生成数据的矩阵乘法。例如，在 TPUv5e 上，在 Gemma 2 9B 模型中寻找前 2% 的隐藏激活值，耗时是相应矩阵乘法的 27 倍。

虽然近似 Top-K 算法提供了一种变通方案，但 Chern 等人（2022）提出的最先进方法（实现为 jax.lax.approx_max_k）仍会产生显著的延迟。他们的两阶段方法将输入划分为桶，从每个桶中选择前 1 个元素，然后对结果进行排序。然而，第二阶段（对选定的元素进行排序）仍然是瓶颈，而第一阶段往往未能充分利用可用的标量/向量计算资源。

方法论
作者提出了对 Chern 等人（2022）算法的广义化，修改了第一阶段，使其从每个桶中选择前 K' 个元素（其中 $1 \le K' \le K$ ），而不是仅限于选择前 1 个。

算法广义化：
- 第一阶段： 将大小为 $N$ 的输入数组划分为 $B$ 个桶。算法不再仅选择最大值，而是从每个桶中选择前 $K'$ 个元素。
- 第二阶段： 合并并排序选定的 $B \times K'$ 个元素，以返回最终的前 $K$ 个元素。
- 实现： 第一阶段使用 Pallas（一种用于 TPU 的 JAX 内核语言）实现为在线更新子程序。它为每个桶维护一个排序好的前 $K'$ 个值及其索引列表。新元素通过类似冒泡排序的传递进行插入，并针对跨桶的向量化进行了优化。作者提供了“未融合”版本（两个独立内核）和“矩阵乘法融合”版本，其中第一阶段与前面的矩阵乘法合并，从而有效地隐藏了第一阶段的开销。
理论分析：
- 作者推导了该广义算法在随机划分下的期望召回率的精确表达式，利用超几何分布对前 $K$ 个元素在桶中的分布进行建模。
- 他们证明，对于固定的期望召回率目标，增加 $K'$ 可以减少桶的数量 $B$ ，使得传递给第二阶段元素总数（ $B \times K'$ ）最小化。
- 至关重要的是，他们推导了当 $K'=1$ 时所需桶数量的新界限。与 Chern 等人（2022）中的界限相比，该界限在理论上严格了 2 倍，即使在原始设置下也能允许使用更少的桶，从而减少需要排序的元素数量。
性能建模：
- 利用脊点分析（ridge-point analysis），作者刻画了加速器的子系统（HBM、VPU、MXU）。他们证明，第一阶段通常受限于内存，直到 $K'$ 达到某个阈值（例如在 TPUv5e 上 $K' \approx 6$ ）。在这个内存受限的范围内，增加 $K'$ 不会显著增加第一阶段的延迟，但会大幅减少第二阶段的工作量。

主要贡献

广义算法： 一种新的两阶段近似 Top-K 算法，每个桶选择前 $K'$ 个元素，在第一阶段计算与第二阶段排序负载之间提供了灵活的权衡。
理论界限：
- 期望召回率作为 $N, K, B$ 和 $K'$ 函数的精确表达式。
- 与 prior work 相比，针对 $K'=1$ 情况所需的桶数量提供了一个严格了 2 倍的界限，改进了参数选择。
实现： 基于 Pallas 在 Cloud TPUv5e 上的高性能实现，包含未融合和矩阵乘法融合两种变体。
自动参数选择： 一种例程，可自动选择 $K'$ 和 $B$ 以满足用户指定的召回率目标，消除了手动调整的需求。

结果

理论效率： 对于固定的召回率目标（例如 95%），增加 $K'$ 可显著减少第二阶段处理的元素数量。例如，从 262,144 个元素中选择前 1024 个，当 $K'=1$ 时需要 16,384 个输出元素，而当 $K'=4$ 时仅需 2,048 个元素（减少了 8 倍）。
延迟改进（未融合）： 在 TPUv5e 上，对于相同的任务，采用 $K'=4$ 的广义算法在 99% 召回率目标下，相比原始 jax.lax.approx_max_k 实现了 11 倍加速（27µs 对比 305µs）。
延迟改进（融合）： 在最大内积搜索（MIPS）工作负载（100 万向量，1024 个查询）中，采用 $K'=4$ 的矩阵乘法融合实现将总运行时间从 594ms（精确 Top-K）降低至 10ms。与稠密基线相比，这代表了 6.7 倍的整体加速，使近似 Top-K 步骤的开销低于矩阵乘法本身的开销。
端到端影响： 在稀疏 Transformer 块基准测试（Gemma 2 9B）中，所提出的方法将稀疏 MLP 块的开销从使用原始算法时比稠密版本慢 2.7 倍，降低为仅比稠密基线多出 5ms 的适度开销。

意义与主张
论文主张，先前近似 Top-K 算法中选择 $K'=1$ 并未充分利用现代加速器上的可用计算资源。通过广义化到 $K' > 1$ ，该算法有效地将瓶颈从排序阶段转移出去，而不会在第一阶段过滤过程中产生显著额外成本，前提是第一阶段仍受限于内存。

作者强调，他们的理论分析为参数选择提供了严谨的基础，并提供了一个用户友好的接口（approx_top_k(array, K, recall_target)），无需手动调整。虽然实证评估主要集中在 TPUv5e 上，但作者断言，该算法、性能建模和分析具有通用性，适用于其他硬件平台（如 GPU）。这项工作表明，原则性的算法广义化，结合硬件感知实现（融合），可以消除 Top-K 在大规模模型训练和推理中的瓶颈地位。