Fast LapSum: Exact Differentiable Top-k at Million Scale

• URL: http://arxiv.org/abs/2608.06912v1
• Date: Fri, 07 Aug 2026 07:44:35 GMT
• Title: Fast LapSum: Exact Differentiable Top-k at Million Scale
• Abstract summary: Fast LapSum is an exact-budget soft top-$k$ primitive whose GPU solver runs in linear time after sorting.<n>We demonstrate Fast LapSum on two demanding applications operating over millions of coordinates inside the training loop.
• Score: 10.106073328854459
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Accelerated Relax-and-Round for Concave Coverage Problems [7.984421247931201]
  We present an accelerated relax-and-round algorithm for concave coverage problems.<n>Our algorithm outperforms approaches that use state-of-the-art LP solvers.
  arXiv  Detail & Related papers  (2026-05-07T20:00:40Z)
• Nearly Minimax Optimal Regret for Learning Linear Mixture Stochastic Shortest Path [80.60592344361073]
  We study the Shortest Path (SSP) problem with a linear mixture transition kernel. An agent repeatedly interacts with a environment and seeks to reach certain goal state while minimizing the cumulative cost. Existing works often assume a strictly positive lower bound of the iteration cost function or an upper bound of the expected length for the optimal policy.
  arXiv  Detail & Related papers  (2024-02-14T07:52:00Z)
• Oblivious Stochastic Composite Optimization [47.48197617884748]
  We show that our algorithms converge without any prior knowledge on the parameters of the problem.<n>All three algorithms work without prior knowledge of the diameter of the feasible set, the Lipschitz constant or smoothness of the objective function.<n>We extend our framework to relative scale and demonstrate the efficiency and robustness of our methods on large-scale semidefinite programs.
  arXiv  Detail & Related papers  (2023-06-30T08:34:29Z)
• Learning distributed representations with efficient SoftMax normalization [3.8673630752805437]
  We propose a linear-time approximation to compute the normalization constants of $rm SoftMax(XYT)$ for embedding vectors with bounded norms.<n>We show on some pre-trained embedding datasets that the proposed estimation method achieves higher or comparable accuracy with competing methods.<n>The proposed algorithm is interpretable and easily adapted to arbitrary embedding problems.
  arXiv  Detail & Related papers  (2023-03-30T15:48:26Z)
• Best Policy Identification in Linear MDPs [70.57916977441262]
  We investigate the problem of best identification in discounted linear Markov+Delta Decision in the fixed confidence setting under a generative model. The lower bound as the solution of an intricate non- optimization program can be used as the starting point to devise such algorithms.
  arXiv  Detail & Related papers  (2022-08-11T04:12:50Z)
• Asynchronous Stochastic Optimization Robust to Arbitrary Delays [54.61797739710608]
  We consider optimization with delayed gradients where, at each time stept$, the algorithm makes an update using a stale computation - d_t$ for arbitrary delay $d_t gradient. Our experiments demonstrate the efficacy and robustness of our algorithm in cases where the delay distribution is skewed or heavy-tailed.
  arXiv  Detail & Related papers  (2021-06-22T15:50:45Z)
• Minimax Optimization with Smooth Algorithmic Adversaries [59.47122537182611]
  We propose a new algorithm for the min-player against smooth algorithms deployed by an adversary. Our algorithm is guaranteed to make monotonic progress having no limit cycles, and to find an appropriate number of gradient ascents.
  arXiv  Detail & Related papers  (2021-06-02T22:03:36Z)
• Kernel methods through the roof: handling billions of points efficiently [94.31450736250918]
  Kernel methods provide an elegant and principled approach to nonparametric learning, but so far could hardly be used in large scale problems. Recent advances have shown the benefits of a number of algorithmic ideas, for example combining optimization, numerical linear algebra and random projections. Here, we push these efforts further to develop and test a solver that takes full advantage of GPU hardware.
  arXiv  Detail & Related papers  (2020-06-18T08:16:25Z)
• Maximizing Determinants under Matroid Constraints [69.25768526213689]
  We study the problem of finding a basis $S$ of $M$ such that $det(sum_i in Sv_i v_i v_itop)$ is maximized. This problem appears in a diverse set of areas such as experimental design, fair allocation of goods, network design, and machine learning.
  arXiv  Detail & Related papers  (2020-04-16T19:16:38Z)
• Sparse Regression at Scale: Branch-and-Bound rooted in First-Order Optimization [6.037383467521294]
  We present a new exact MIP framework for $ell_0$ regularized regression. Our framework can scale to $p sim 107$, achieving speedups of at least $5000$x. We open source the implementation through our toolkit L0BnB.
  arXiv  Detail & Related papers  (2020-04-13T18:45:29Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.