提出可微分的精确顶K方法,百万级规模下毫秒级计算。
Fast LapSum: Exact Differentiable Top-k at Million Scale

- 基于线性时间阈值与解析雅可比,实现精确预算的可微顶K
- 处理百万、千万、亿级数据仅需0.41/1.15/5.23毫秒
- 适用于稀疏路由、对抗样本生成等大规模训练场景
顶K操作是现代稀疏计算的核心,用于标记路由、专家激活、内存选择和注意力剪枝。但标准硬顶K无法传递梯度,而现有连续松弛方法在大规模模型中仍过于昂贵。本文提出Fast LapSum,一种精确预算的可微顶K原语,其GPU求解器在排序后为线性时间。不同于此前线性时间方法(如DFTopK)放松归一化约束,Fast LapSum是首个在保持精确选择总量 $k$ 的同时完全可微端到端的方法。求解器结合线性时间阈值计算与解析向量-雅可比乘积,并在极端规模下采用概率夹钳策略,仅对核噪声得分中不确定的中间区域排序。开销几乎可忽略:处理 $10^6$、$10^7$、$10^8$ 个分数分别耗时 0.41、1.15、5.23 毫秒。该方法使精确软顶K在稀疏路由、检索和大规模优化中成为可行方案。我们在两个高要求应用中验证:在训练循环内处理数百万坐标,生成像素级稀疏对抗样本(精确软预算约图像像素的 0.02%),相比最先进方法提速一个数量级;并从零开始训练全可微稀疏图像编码器。
原文摘要 · Abstract (English)
The top-$k$ operation is a fundamental building block of modern sparse computation, enabling token routing, expert activation, memory selection, and attention pruning. Yet standard hard top-$k$ blocks gradients, while existing continuous (soft) relaxations remain too costly for large-scale models. We introduce Fast LapSum, an exact-budget soft top-$k$ primitive whose GPU solver runs in linear time after sorting. Unlike prior linear-time methods such as DFTopK, which relax the normalization constraint, Fast LapSum is, to our knowledge, the first method to preserve an exact selection mass of $k$ while remaining fully differentiable end-to-end. Our solver combines a linear-time threshold computation with an analytical vector--Jacobian product, and for extreme scales employs probabilistic bracketing to sort only the uncertain middle band of kernel-noised scores. The resulting overhead is almost negligible: the solver processes $10^6$, $10^7$, and $10^8$ scores in $0.41$, $1.15$, and $5.23$\,ms, respectively. This makes exact soft top-$k$ practical for sparse routing, retrieval, and large-scale optimization. We demonstrate Fast LapSum on two demanding applications operating over millions of coordinates inside the training loop: generating megapixel sparse adversarial examples with an exact soft budget of ${\sim}0.02\%$ of an image's pixels, achieving an order-of-magnitude speedup over state-of-the-art methods, and training a fully differentiable sparse image coder from scratch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。