🤖 AI

Fast LapSum: Exact Differentiable Top-k at Million Scale

The paper introduces Fast LapSum, an exact, differentiable soft top- $k$ primitive that preserves a precise selection mass of $k$ while running in linear time on GPUs, enabling efficient million-scale sparse computation for applications like adversarial example generation and differentiable image coding.

Original authors: Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

Published 2026-08-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive digital library where millions of books are being scanned every second. To make sense of this flood of information, the library's AI needs to decide which few books are the most important to read right now. In the world of artificial intelligence, this is called "top-k selection": picking the best k items from a huge list. Usually, the AI does this by being a strict librarian who picks the top books and ignores the rest completely. This is great for speed, but it's terrible for learning because the AI can't figure out how to get better; it's like trying to learn to drive by only looking at the road when you're already in the right lane, with no way to adjust the steering wheel.

To fix this, scientists invented "soft" versions of this selection. Instead of a hard "yes or no," the AI gives a "maybe" score to every book, allowing it to learn from the mistakes. But here's the catch: these soft versions are often so slow and computationally heavy that they crash the system when the library gets too big. They are like trying to sort a million books by hand while the library is on fire. The big question for researchers has been: Can we have a librarian that is both gentle enough to learn (differentiable) and fast enough to handle millions of books without breaking a sweat?

This is where the new paper, "Fast LapSum," steps in. The authors, a team from Poland, have built a new tool that acts like a super-efficient, mathematically perfect librarian. They created a method called Fast LapSum that allows an AI to pick the top items from a list of millions while still being able to learn from the process. Unlike previous methods that either gave up on being perfectly accurate to gain speed, or were too slow to be useful, Fast LapSum manages to do both. It finds the exact right number of items to pick (the "budget") and calculates the perfect "maybe" scores for them in the blink of an eye.

The secret sauce is a clever trick involving a "blurred" view of the scores. Imagine the scores aren't sharp points but fuzzy clouds. The AI needs to draw a line through these clouds so that the total amount of "cloud" above the line equals exactly the number of books it's allowed to pick. Old methods tried to find this line by guessing and checking over and over, which took forever. Fast LapSum, however, uses a special mathematical formula (based on something called the Laplace distribution) that lets it calculate the line instantly after a single sort.

For really huge lists—like a million or even a hundred million scores—the authors added a second trick called "probabilistic bracketing." Instead of sorting the entire million-item list, which is like trying to organize a stadium full of people, the system takes a quick sample to guess where the line probably is. It then only sorts the small group of people standing right near that line. This keeps the process incredibly fast, taking just a few milliseconds even for massive datasets.

The paper proves this works by testing it on two very difficult tasks. First, they used it to create "adversarial examples," which are images that look normal to humans but trick AI classifiers. They managed to change an image so slightly—altering only about 0.02% of the pixels (roughly 600 pixels out of 3.3 million)—that the AI misidentified a picture of a tiger. This was done much faster and with less "damage" to the image than previous methods. Second, they built a differentiable image coder from scratch, a system that compresses images by selecting only the most important parts to keep. In both cases, Fast LapSum acted as the engine, handling millions of decisions per second without slowing down the learning process.

The authors show that this method is not just a theoretical idea but a practical tool that runs in milliseconds on standard computer chips. They compared their work to other recent attempts, like one called DFTopK, and found that while those methods are fast, they sacrifice the exactness of the selection (the total number of items picked drifts away from the target). Fast LapSum, they argue, is the first to keep the selection perfectly exact while remaining fast enough for real-world, large-scale AI systems. It turns a slow, expensive bottleneck into a smooth, fast operation, allowing AI to be both smart and efficient.

Technical Summary: Fast LapSum: Exact Differentiable Top-k at Million Scale

1. Problem Statement

The top-k operation is a fundamental primitive in modern sparse computation, essential for token routing in Mixture-of-Experts (MoE) models, memory selection in long-context architectures, and retrieval systems. However, standard hard top-k selection is non-differentiable, blocking gradient flow across the decision boundary. Conversely, existing continuous (soft) relaxations (e.g., NeuralSort, SoftSort, Sinkhorn-OT) often introduce prohibitive computational costs ( $O(n^2)$ or iterative complexity) or fail to preserve the exact selection budget (the sum of selected probabilities equals $k$ ).

Current differentiable top-k methods face a trade-off: they are either too expensive to run within the training loop of large-scale models or they relax the normalization constraint ( $\sum p_i = k$ ), leading to uncontrolled sparsity. The paper argues that a differentiable sparsifier must be a "systems primitive" that maintains the exact budget of hard top-k while remaining computationally cheap enough for million-scale vectors.

2. Methodology

2.1 The LapSum Operator

The authors build upon the LapSum operator, which models top-k selection probabilistically. Instead of a deterministic binary mask, LapSum treats scores $r_i$ as perturbed by Laplace noise $\epsilon \sim \text{Laplace}(0, t)$ . The selection probability $p_i$ is defined as the probability that a perturbed score exceeds a scalar threshold $b$ :
$p_i = \sigma\left(\frac{r_i - b}{t}\right)$
where $\sigma$ is the Laplace Cumulative Distribution Function (CDF). The threshold $b$ is uniquely determined by the constraint that the total mass equals the budget $k$ :
$\sum_{i=1}^n p_i = k$
Unlike prior methods that approximate this threshold, LapSum solves for $b$ exactly. The backward pass is analytical, utilizing a vector-Jacobian product (VJP) that avoids biased surrogates like straight-through estimators.

2.2 Fast LapSum: The GPU Solver

The core contribution is Fast LapSum, a GPU-native solver that computes the exact threshold $b$ in linear time after sorting.

Full-Sort Solver: For vectors where sorting is feasible, the method computes a "threshold curve" using a single stable prefix/suffix scan over the sorted scores. The threshold $b$ is found via an index lookup and a closed-form root calculation within the specific interval where the budget equation balances. This replaces the iterative bisection used in previous implementations, reducing post-sorting work to $O(n)$ .
Probabilistic Bracketing (Million-Scale): For extreme scales (e.g., $n > 10^6$ $n > 1 0^{6}$ ) where sorting the entire vector becomes the bottleneck, the authors introduce a probabilistic bracketing pipeline:
1. Localization: A small sample of kernel-noised scores is drawn to estimate the quantile corresponding to the budget. Using binomial order statistics, a narrow interval $[a_L, a_R)$ is identified that contains the true threshold $b$ with high probability.
2. Streaming and Filtering: The full vector is streamed once. Scores outside the bracket are aggregated into stable tail sums ( $\tilde{S}_L, \tilde{S}_R$ ) based on the Laplace CDF properties. Only the "uncertain middle band" of scores within $[a_L, a_R)$ is retained.
3. Exact Solution: The exact LapSum solver is applied only to this reduced middle band.
4. Verification: A certification step checks if the bracket holds; if not, it widens iteratively. This ensures the final result satisfies the exact budget constraint $\sum p_i = k$ .

2.3 Relation to DFTopK

The paper distinguishes Fast LapSum from the concurrent DFTopK method. DFTopK achieves linear time by fixing the threshold at the midpoint between the $k$ -th and $(k+1)$ -th order statistics, effectively relaxing the normalization constraint. The authors prove that DFTopK is the zero-temperature limit of LapSum. However, at finite temperatures, Fast LapSum preserves the exact budget, whereas DFTopK allows the active mass to drift (e.g., resulting in a mass of $\approx 3.46k$ instead of $k$ ).

3. Key Contributions

Fast LapSum Operator: The first GPU solver to achieve linear post-sorting complexity while strictly preserving the unrelaxed constraint $\sum p_i = k$ . It eliminates the need for biased surrogates.
Probabilistic Bracketing Pipeline: A method for million-scale vectors that localizes the threshold via sampling, sorts only the uncertain middle band, and verifies the exact solution, maintaining expected linear time complexity.
Scalability and Efficiency: Demonstrated ability to process vectors up to $10^8$ scores in low-millisecond times (e.g., 5.23 ms for $10^8$ scores), orders of magnitude faster than previous exact implementations.

4. Experimental Results

4.1 Performance Benchmarks

Speed: Fast LapSum's full-sort branch is nearly two orders of magnitude faster than the per-query bisection baseline. The bracketed branch achieves speedups of up to 19 $\times$ over full sorting at $10^8$ scores and 1200 $\times$ over the official LapSum release at $3.2 \times 10^7$ scores.
Memory: Unlike dense sorting relaxations (NeuralSort, SoftSort) which run out of memory (OOM) around $n=10^5$ , Fast LapSum operates efficiently at $n=10^8$ .
Accuracy: Fast LapSum maintains exact budget normalization (relative error $< 10^{-5}$ ) across all scales. In contrast, DFTopK exhibits a relative error of $\approx 2.46$ (mass $\approx 3.46k$ ) due to its relaxed constraint.

4.2 Applications

The authors demonstrate the operator in two demanding, million-scale applications:

Differentiable Megapixel Image Coder: A fully differentiable sparse image coder trained from scratch. The operator manages coordinate selection over millions of scores within a rate-distortion loop, successfully assembling an image from sparse atoms without dominating the computational cost.
Sparse Adversarial Attacks: The operator is used to generate sparse adversarial examples on native-resolution images (DIV2K) with a ConvNeXt-V2 classifier.
- Efficiency: The attack achieves a 100% success rate with a median budget of 199.8 fully-modified-pixel equivalents.
- Comparison: This is significantly more efficient than baselines: $\sigma$ -zero requires 3,101, Sparse-PGD requires 16,134, and PGD-L0 requires 74,553.
- Precision: The attack concentrates perturbations on class-defining textures (e.g., the eye of a tiger) rather than scattered noise, using only $\sim 0.02\%$ of the image pixels.

5. Significance and Claims

The paper claims that Fast LapSum resolves the long-standing trade-off between differentiability, exact budget constraints, and computational efficiency. By treating differentiable top-k as a systems primitive, the method enables:

End-to-End Training: The operator is fast enough to be embedded directly within the training loop of large sparse models (routing, retrieval, pruning) without becoming the computational bottleneck.
Exact Sparsity Control: Unlike soft relaxations that drift in mass, Fast LapSum guarantees that the number of active experts, memory entries, or retrieval candidates strictly adheres to the specified budget $k$ .
Practicality at Scale: The method makes exact soft top-k practical for real-world tasks involving millions of coordinates, such as high-resolution image processing and large-scale retrieval, where previous exact methods were infeasible.

The authors note limitations: the closed-form solution is specific to the Laplace kernel, and the linear scaling is expected rather than worst-case (pathological score distributions could trigger extra verification passes). However, the empirical results suggest the operator is a negligible cost at scale, enabling a new class of exact-budget differentiable systems.