🤖 AI

Fast LapSum: Exact Differentiable Top-k at Million Scale

该论文介绍了 Fast LapSum，一种精确、可微的 soft top- $k$ 原语，它在保持精确的 $k$ 选择质量的同时，能在 GPU 上以线性时间运行，从而为对抗样本生成和可微图像编码等应用实现高效的百万级规模稀疏计算。

原作者： Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

发布于 2026-08-10

📖 1 分钟阅读☕ 轻松阅读

原作者： Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

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

想象一下，你正在经营着一个规模巨大的数字图书馆，每秒钟都有数百万本书籍正在被扫描。为了从这些如洪水般涌入的信息中理出头绪，图书馆的人工智能需要决定现在最值得阅读的几本书。在人工智能的世界里，这被称为“top-k 选择”：从庞大的列表中挑选出前 $k$ 个项目。通常情况下，人工智能通过扮演一名严格的图书管理员来完成这项工作，只挑选出顶尖的书籍并完全忽略其余部分。这种方式虽然速度很快，但对于学习过程来说却很糟糕，因为人工智能无法弄清楚如何变得更好；这就像是试图通过只看正确车道上的路来学习驾驶，却没有任何调整方向盘的方法。

为了解决这个问题，科学家们发明了“软性”版本的选择方法。它不再是硬性的“是或否”，而是给每一本书一个“可能”的分数，从而允许人工智能从错误中学习。但问题在于，这些软性版本通常计算量巨大且极其缓慢，当图书馆规模变大时会导致系统崩溃。它们就像是在图书馆失火时，试图靠手工对一百万本书进行分类一样。一个核心问题一直是：能否创造出一位既足够温柔以进行学习（可微），又足够快速以处理数百万本书籍而不精疲力竭的图书管理员？

这正是新论文《Fast LapSum》介入的地方。作者们——来自波兰的一个团队——构建了一个全新的工具，它充当了一位超级高效、在数学上近乎完美的图书管理员。他们创建了一种名为 Fast LapSum 的方法，允许人工智能从数百万个项目中挑选出前 $k$ 个项目，同时仍能从这个过程中学习。与以往那些要么为了追求速度而放弃完美准确性，要么因速度太慢而无法投入实际使用的各种方法不同，Fast LapSum 兼顾了两者。它能找到要挑选的精确数量（即“预算”），并在眨眼之间为这些项目计算出完美的“可能”分数。

其背后的秘诀是一个涉及分数“模糊化”视图的巧妙技巧。想象一下，这些分数不是锐利的点，而是模糊的云团。人工智能需要在这些云团中画一条线，使得线以上的“云团”总量恰好等于它被允许挑选的书籍数量。旧的方法试图通过反复的猜测和尝试来寻找这条线，这非常耗时。然而，Fast Lapsum 使用了一个基于拉普拉斯分布（Laplace distribution）的特殊数学公式，这让它在仅进行一次排序后就能瞬间计算出这条线。

对于规模极大的列表——比如一百万甚至一亿个分数——作者们加入了第二个技巧，称为“概率分级法”（probabilistic bracketing）。与其对整个百万级的列表进行排序（这就像是在组织一个坐满了人的体育场），该系统会先进行快速采样，以推测这条线大概位于何处。然后，它只对紧邻该线的一小部分人群进行排序。这使得整个过程极其迅速，即使面对海量数据集，也仅需几毫秒即可完成。

论文通过在两个极具挑战性的任务上进行测试，证明了这一方法的有效性。首先，他们利用它创建了“对抗样本”，即那些看起来对人类很正常但能欺骗 AI 分类器的图像。他们成功地微调了一张图像，仅改变了约 0.02% 的像素（在 330 万个像素中大约只有 600 个像素），就让 AI 误认了一张老虎的照片。与之前的方法相比，这一过程更快，且对图像造成的“损伤”更小。其次，他们从零开始构建了一个可微图像编码器，这是一个通过仅保留最重要的部分来进行图像压缩的系统。在这两个案例中，Fast LapSum 都充当了引擎，在每秒处理数百万次决策的同时，并未减缓学习进程。

作者展示了这种方法不仅是一个理论构想，更是一个能在标准计算机芯片上以毫秒级运行的实用工具。他们将自己的工作与名为 DFTopK 的近期研究进行了对比，发现虽然那些方法很快，但它们牺牲了选择的精确性（即所选项目的总数会偏离目标值）。Fast LapSum 则认为，它是第一个在保持极速的同时，仍能保持选择过程完美精确的方法。它将一个缓慢、昂贵的瓶颈转化为了一个流畅、快速的操作，让人工智能既聪明又高效。

技术摘要：Fast LapSum：百万规模下的精确可微 Top-k

1. 问题陈述

Top-k 操作是现代稀疏计算中的一个基本原语，对于混合专家（MoE）模型中的 Token 路由、长上下文架构中的内存选择以及检索系统至关重要。然而，标准的硬 top-k 选择是不可微的，这阻断了跨越决策边界的梯度流。相反，现有的连续（软）松弛方法（例如 NeuralSort、SoftSort、Sinkhorn-OT）通常会引入极高的计算成本（ $O(n^2)$ 或迭代复杂度），或者无法保持精确的选择预算（即所选概率之和等于 $k$ ）。

目前的微分 top-k 方法面临着一种权衡：它们要么在大型模型的训练循环中运行得过于昂贵，要么无法维持归一化约束（ $\sum p_i = k$ ），从而导致不可控的稀疏性。本文认为，一个可微的稀疏化算子必须是一个“系统原语”，既能保持硬 top-k 的精确预算，又能对百万规模的向量保持足够低的计算成本。

2. 方法论

2.1 LapSum 算子

作者基于 LapSum 算子构建，该算子通过概率方式对 top-k 选择进行建模。LapSum 不使用确定性的二值掩码，而是将分数 $r_i$ 视为受拉普拉斯噪声 $\epsilon \sim \text{Laplace}(0, t)$ 扰动的变量。选择概率 $p_i$ 定义为扰动后的分数超过标量阈值 $b$ 的概率：
$p_i = \sigma\left(\frac{r_i - b}{t}\right)$
其中 $\sigma$ 是拉普拉斯累积分布函数（CDF）。阈值 $b$ 是根据总质量等于预算 $k$ 的约束条件唯一确定的：
$\sum_{i=1}^n p_i = k$
与以往近似该阈值的方法不同，LapSum 精确求解 $b$ 。其反向传播是解析式的，利用向量-雅可比乘积（VJP）避免了像直通估计器（straight-through estimators）那样具有偏差的替代方案。

2.2 Fast LapSum：GPU 求解器

核心贡献是 Fast LapSum，这是一个能够在排序后实现线性时间求解精确阈值 $b$ 的 GPU 原生求解器。

全排序求解器（Full-Sort Solver）： 对于排序可行的向量，该方法通过对排序后的分数进行单次稳定的前缀/后缀扫描来计算“阈值曲线”。阈值 $b$ 通过索引查找以及在预算方程平衡的特定区间内的闭式根计算得出。这取代了以往实现中使用的迭代二分法，将排序后的工作量降低至 $O(n)$ 。
概率区间化（百万规模）（Probabilistic Bracketing）： 对于极端规模（例如 $n > 10^6$ $n > 1 0^{6}$ ）下排序整个向量成为瓶颈的情况，作者引入了一个概率区间化流水线：
1. 定位（Localization）： 抽取少量经过核噪声处理的分数样本，以估计对应于预算的分位数。利用二项分布序统计量，识别出一个以高概率包含真实阈值 $b$ 的狭窄区间 $[a_L, a_R)$ 。
2. 流式处理与过滤（Streaming and Filtering）： 对全量向量进行一次流式处理。根据拉普拉斯 CDF 的性质，将区间外的分数聚合为稳定的尾部和（ $\tilde{S}_L, \tilde{S}_R$ ）。仅保留位于 $[a_L, a_R)$ 范围内的“不确定中间带”分数。
3. 精确求解（Exact Solution）： 仅对这个缩减后的中间带应用精确的 LapSum 求解器。
4. 验证（Verification）： 通过一个认证步骤检查区间是否成立；如果不成立，则迭代扩大区间。这确保了最终结果满足精确预算约束 $\sum p_i = k$ 。

2.3 与 DFTopK 的关系

本文将 Fast Lap-Sum 与同时期的 DFTopK 方法进行了区分。DFTopK 通过将阈值固定在第 $k$ 和第 $(k+1)$ 个序统计量之间来实现线性时间复杂度，这实际上放宽了归一化约束。作者证明了 DFTopK 是 LapSum 在零温度极限下的表现。然而，在有限温度下，Fast LapSum 保持了精确的预算，而 DFTopK 则允许活跃质量发生漂移（例如，导致活跃质量约为 $3.46k$ 而不是 $k$ ）。

3. 核心贡献

Fast LapSum 算子： 第一个实现线性排序复杂度且严格保持非松弛约束 $\sum p_i = k$ 的 GPU 求解器。它消除了对有偏差替代方案的需求。
概率区间化流水线： 一种针对百万规模向量的方法，通过采样定位阈值，仅对不确定的中间带进行排序，并验证精确解，从而维持期望的线性时间复杂度。
可扩展性与效率： 证明了处理高达 $10^8$ 个分数的向量时仅需低毫秒级时间（例如， $10^8$ 个分数耗时 5.23 ms），其速度比之前的精确实现高出数个数量级。

4. 实验结果

4.1 性能基准测试

速度： Fast LapSum 的全排序分支比每查询一次二分法的基准快近两个数量级。区间化分支在 $10^8$ 个分数时比全排序快达 19 倍，在 $3.2 \times 10^7$ 个分数时比官方 LapSum 发布版本快 1200 倍。
内存： 不同于会在 $n=10^5$ 左右出现显存溢出（OOM）的稠密排序松弛方法（如 NeuralSort, SoftSort），Fast LapSum 在 $n=10^8$ 时仍能高效运行。
准确度： Fast LapSum 在所有规模下都能保持精确的预算归一化（相对误差 $< 10^{-5}$ ）。相比之下，由于其松弛的约束，DFTopK 的相对误差约为 $2.46 $（质量约为$ 3.46k$）。

4.2 应用场景

作者在两个高要求的百万规模应用中展示了该算子：

可微兆像素图像编码器（Differentiable Megapixel Image Coder）： 一个从头开始训练的全可微稀疏图像编码器。该算子在率失真循环中管理数百万个分数的坐标选择，在不主导计算成本的前提下，成功地从稀疏原子中组装图像。
稀疏对抗攻击（Sparse Adversarial Attacks）： 该算子用于在原生分辨率图像（DIV2K）上针对 ConvNeXt-V2 分类器生成稀疏对抗样本。
- 效率： 该攻击实现了 100% 的成功率，其中值预算为 199.8 个完全修改像素等效值。
- 对比： 这比基准方法显著更高效： $\sigma$ -zero 需要 3,101 个，Sparse-PGD 需要 16,134 个，而 PGD-L0 需要 74,553 个。
- 精度： 攻击将扰动集中在定义类别的纹理上（例如老虎的眼睛），而非散布的噪声，仅使用了约 $0.02\%$ 的图像像素。

5. 意义与主张

本文声称 Fast LapSum 解决了长期以来在可微性、精确预算约束和计算效率之间的权衡问题。通过将可微 top-k 处理为系统原语，该方法实现了：

端到端训练： 该算子足够快，可以直接嵌入到大型稀疏模型的训练循环中（如路由、检索、剪枝），而不会成为计算瓶颈。
精确的稀疏性控制： 不同于会出现质量漂移的软松弛方法，Fast LapSum 保证了活跃专家、内存条目或检索候选者的数量严格遵循指定的预算 $k$ 。
大规模实用性： 该方法使精确的软 top-k 在涉及数百万坐标的实际任务（如高分辨率图像处理和大规模检索）中变得可行，而此前的精确方法在这些任务中是不可行的。

作者指出了一些局限性：闭式解是针对拉普拉斯核特有的，且线性缩放是期望值而非最坏情况（病态的分数分布可能会触发额外的验证过程）。然而，实验结果表明，该算子在大规模场景下成本极低，从而开启了一类新的精确预算可微系统。