跳到正文
arXiv cs.AI· Jakub Antczak, Joanna Wojciechowicz, Kamil Ksi\k{a}\.zek, Marcin Mazur, {\L}ukasz Struski, Jacek Tabor·· 6 天前AI 评分26

Fast LapSum:百万级规模下精确可微 Top-k 求解器

Fast LapSum: Exact Differentiable Top-$k$ at Million Scale

AI 导读

Fast LapSum 是一种可扩展的 LapSum 软 top-k 求解器,在保持精确选择质量 k 的同时支持端到端可微。其认证 GPU 实现处理 10^6、10^7、10^8 个分数中位耗时分别为 0.92、1.39、7.24 ms。在百万像素稀疏对抗样本上较 SOTA 方法实现一个数量级加速,并用于 3D Gaussian splatting 减少高斯数量。

正文 · AI 翻译

查看 PDF HTML(实验性)

摘要:选取前 $k$ 个元素是在大规模模型和优化问题中诱导稀疏性的基本操作,可实现稳健的专家激活、token 路由或注意力剪枝。然而,硬 top-$k$ 不可微,而现有的可微替代方案随着坐标数量的增长而变得愈发昂贵。我们提出 Fast LapSum,一种用于 LapSum 软 top-$k$ 公式的可扩展求解器,它保持精确的选择质量 $k$,同时支持端到端微分。在 Fast LapSum 中,我们使用概率括界来降低排序成本,通过基于核噪声样本的二项式顺序统计量,将排序限制在阈值附近的狭窄分数带内。一次认证过程以额外一次线性遍历的代价将概率定位升级为已验证定位,并带有覆盖最坏情况的完整排序回退。我们经过认证的 GPU 实现分别以 $0.92$、$1.39$ 和 $7.24$ 毫秒的中位时间处理 $10^6$、$10^7$ 和 $10^8$ 个分数,使精确预算的软 top-$k$ 在百万级优化循环中变得实用。我们在两个应用中展示了这一能力:使用少量初始图像像素的百万像素稀疏对抗样本,其中 Fast LapSum 相比最先进方法实现了一个数量级的加速;以及 3D 高斯泼溅。在后者中,我们使用 Fast LapSum 将自适应密度控制产生的高斯数量减少到显著更少的数量,同时保持几乎相同的渲染质量并大幅减少计算量。这些结果表明,精确预算的可微 top-$k$ 可以融入实际的百万级优化流水线。
学科: 人工智能(cs.AI)
引用为: arXiv:2608.06912 [cs.AI]
  (或使用 arXiv:2608.06912v3 [cs.AI] 查看此版本)
  https://doi.org/10.48550/arXiv.2608.06912

由 DataCite 颁发的 arXiv DOI

提交历史

来自:Jakub Antczak [查看电子邮件]
[v1] 2026年8月7日星期五 07:44:35 UTC(2,636 KB)
[v2] 2026年9月30日星期三 11:24:22 UTC(4,513 KB)
[v3] 2026年10月8日星期四 18:50:00 UTC(4,519 KB)

来源:arXiv cs.AI · arxiv.org