Fast LapSum: Exact Differentiable Top-k at Million Scale

2026-08-07

AI Summary

Hard top-k operations are non-differentiable, and existing soft approximations suffer from either high computational cost or insufficient selection accuracy in large-scale settings. This work proposes Fast LapSum, the first soft top-k operator that simultaneously achieves exact k-selection, full differentiability, and linear time complexity. Its key innovations include a linear-time thresholding algorithm based on Laplace noise, an analytic vector-Jacobian product, and a probabilistic bracketing strategy for sorting uncertain score segments. Fast LapSum processes millions to hundreds of millions of scores in only 0.41–5.23 milliseconds, accelerates the generation of ultra-sparse (0.02% pixel) adversarial examples by an order of magnitude over prior methods, and enables, for the first time, end-to-end training of a fully differentiable sparse image encoder.

Abstract

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.

Problem

• top-k
• differentiable
• sparse computation
• large-scale optimization
• soft relaxation

Innovation

• differentiable top-k
• exact budget
• linear-time solver
• sparse computation
• probabilistic bracketing