🤖 machine learning

A Faster Generalized Two-Stage Approximate Top-K

This paper generalizes a two-stage approximate Top-K algorithm by selecting top- $K'$ elements per partition instead of just the top-1, providing a tighter theoretical recall bound and demonstrating an order-of-magnitude speedup on Cloud TPUv5e while maintaining the same expected recall.

Original authors: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

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 the manager of a massive library with millions of books (data). Every day, you need to find the Top-K most popular books (the K largest numbers) to recommend to visitors.

In the world of computer chips (specifically the ones used to train giant AI models), finding these "most popular" items is surprisingly slow and expensive. It's like trying to find the top 100 books by reading every single one of them, one by one, even though your library is designed to do math on huge stacks of books all at once.

Here is the simple breakdown of what this paper does to fix that problem.

The Old Way: The "One-at-a-Time" Filter

A previous method (by Chern et al., 2022) tried to speed this up using a two-step process:

The Split: Imagine dividing your library into 100 different rooms (buckets).
The First Scan: In each room, a helper picks only the single most popular book and brings it to the front desk.
The Final Sort: The manager then looks at just those 100 books (one from each room) and picks the top 100 overall.

The Problem: This method was too cautious. By only picking the one best book from each room, it often missed the second or third best books that were hiding in the same room. To make sure it didn't miss anything, they had to use many rooms (buckets), which meant the manager still had to sort a huge pile of books at the end. It was still too slow.

The New Idea: The "Top-K" Filter

The authors of this paper realized that the computer chips have extra power they weren't using. They proposed a smarter version of the first step:

Instead of picking just the #1 book from each room, the helper now picks the Top-K' books (for example, the top 4) from each room.

Why is this better?

Fewer Rooms Needed: Because the helper is grabbing more books from each room, you don't need as many rooms to ensure you catch all the popular books.
Less Sorting: Even though the helper grabs more books per room, the total number of books sent to the manager for the final sort is actually much smaller.
The Result: The manager has a tiny pile to sort instead of a mountain.

The "Magic" of the Hardware

The paper explains that modern computer chips (like Google's TPU) are like giant factories with different workstations:

The Matrix Unit (MXU): A super-fast factory that does heavy math (multiplication) but is bad at sorting.
The Vector Unit (VPU): A smaller, slower workstation that is good at sorting and picking winners.

The old method wasted the VPU's time. The new method uses the VPU to grab the "Top-K'" books while the MXU is busy doing math. It's like having a worker grab the best items from a conveyor belt while the machine is still running, so there's no waiting time.

The Results: Speeding Up the AI

The authors tested this on a Google TPU chip:

The Old Way: Finding the top books took a long time, often slower than the math that created the list in the first place.
The New Way: By grabbing the "Top 4" from each bucket instead of just the "Top 1," they reduced the work for the final sort by 7 times on average.
The Fusion: They even managed to combine the "picking" step with the "math" step so they happen at the exact same time.

The Bottom Line:
In a real-world test (finding the top 2% of data in a large AI model), their new method made the process 24 times faster than the previous standard. This means the AI model can train and run much faster without losing accuracy.

Summary Analogy

Old Method: You have 1,000 teams. Each team sends you their best player. You then have to interview 1,000 players to find the top 100.
New Method: You have fewer teams (say, 250). Each team sends you their top 4 players. You only have to interview 1,000 players (250 teams × 4 players), but because you got more options from each team, you are just as likely to find the true best players, and you do it much faster because you organized the teams better.

The paper proves mathematically that this "Top-K'" approach is not just a guess; it is a guaranteed way to get the same quality of results with significantly less work.

Technical Summary: A Faster Generalized Two-Stage Approximate Top-K

Problem Statement
The identification of the top- $K$ elements in an array is a fundamental operation in many machine learning algorithms, including Maximum Inner Product Search (MIPS), K-Nearest-Neighbors (KNN), and the optimization of large foundation models (e.g., for sparsity in classification logits, MLP blocks, attention mechanisms, and KVCache compression). On modern accelerators like TPUs and GPUs, computing exact Top- $K$ is often a performance bottleneck. These devices are optimized for dense matrix multiplications (MXUs/TensorCores), whereas Top- $K$ operations typically rely on scalar/vector units (VPUs/CUDA cores) and can be significantly slower than the matrix multiplications that generate the data. For instance, on TPUv5e, finding the top-2% of hidden activations in a Gemma 2 9B model takes 27 times longer than the corresponding matrix multiplication.

While approximate Top- $K$ algorithms offer a workaround, the state-of-the-art method by Chern et al. (2022), implemented as jax.lax.approx_max_k, still incurs significant latency. Their two-stage approach partitions the input into buckets, selects the top-1 element from each bucket, and sorts the results. However, the second stage (sorting the selected elements) remains a bottleneck, and the first stage often underutilizes available scalar/vector compute resources.

Methodology
The authors propose a generalization of the Chern et al. (2022) algorithm that modifies the first stage to select the top- $K'$ elements from each bucket (where $1 \le K' \le K$ ) rather than restricting selection to the top-1.

Algorithmic Generalization:
- Stage 1: The input array of size $N$ is partitioned into $B$ buckets. Instead of selecting only the maximum, the algorithm selects the top- $K'$ elements from each bucket.
- Stage 2: The $B \times K'$ selected elements are merged and sorted to return the final top- $K$ elements.
- Implementation: The first stage is implemented as an online update subroutine using Pallas (a JAX kernel language for TPUs). It maintains a sorted list of top- $K'$ values and indices for each bucket. New elements are inserted using a bubble-sort-like pass, optimized for vectorization across buckets. The authors provide both an "unfused" version (two separate kernels) and a "matmul-fused" version where the first stage is merged with the preceding matrix multiplication, effectively hiding the first-stage cost.
Theoretical Analysis:
- The authors derive an exact expression for the expected recall of this generalized algorithm under random partitioning, modeling the distribution of top- $K$ elements into buckets using the hypergeometric distribution.
- They prove that for a fixed expected recall target, increasing $K'$ allows for a reduction in the number of buckets $B$ such that the total number of elements passed to the second stage ( $B \times K'$ ) is minimized.
- Crucially, they derive a new bound for the number of buckets required when $K'=1$ . This bound is provably tighter by a factor of 2 compared to the bound in Chern et al. (2022), allowing for fewer buckets and thus fewer elements to sort even in the original setting.
Performance Modeling:
- Using ridge-point analysis, the authors characterize the accelerator's subsystems (HBM, VPU, MXU). They demonstrate that the first stage is typically memory-bound until $K'$ reaches a certain threshold (e.g., $K' \approx 6$ on TPUv5e). Within this memory-bound regime, increasing $K'$ does not significantly increase the latency of the first stage but drastically reduces the workload of the second stage.

Key Contributions

Generalized Algorithm: A new two-stage approximate Top- $K$ algorithm that selects top- $K'$ elements per bucket, offering a flexible trade-off between first-stage compute and second-stage sorting load.
Theoretical Bounds:
- An exact expression for expected recall as a function of $N, K, B,$ and $K'$ .
- A provably tighter bound (factor of 2) for the number of buckets required in the $K'=1$ case compared to prior work, improving parameter selection.
Implementation: A high-performance implementation on Cloud TPUv5e using Pallas, featuring both unfused and matmul-fused variants.
Automatic Parameter Selection: A routine that automatically selects $K'$ and $B$ to meet a user-specified recall target, eliminating the need for manual tuning.

Results

Theoretical Efficiency: For a fixed recall target (e.g., 95%), increasing $K'$ significantly reduces the number of elements processed in the second stage. For example, selecting top-1024 from 262,144 elements requires 16,384 output elements with $K'=1$ , but only 2,048 elements with $K'=4$ (an 8x reduction).
Latency Improvements (Unfused): On TPUv5e, the generalized algorithm with $K'=4$ achieves an 11x speedup over the original jax.lax.approx_max_k (27µs vs 305µs) for a 99% recall target on the same task.
Latency Improvements (Fused): In a Maximum Inner Product Search (MIPS) workload (1M vectors, 1024 queries), the matmul-fused implementation with $K'=4$ reduces the total runtime from 594ms (exact Top-K) to 10ms. This represents a 6.7x overall speedup compared to the dense baseline, bringing the approximate Top-K step below the cost of the matrix multiplication itself.
End-to-End Impact: In a sparse transformer block benchmark (Gemma 2 9B), the proposed method reduces the overhead of the sparse MLP block from 2.7x slower than dense (using the original algorithm) to only a modest +5ms overhead relative to the dense baseline.

Significance and Claims
The paper claims that the choice of $K'=1$ in previous approximate Top-K algorithms does not fully utilize available compute resources on modern accelerators. By generalizing to $K' > 1$ , the algorithm effectively shifts the bottleneck away from the sorting stage without incurring significant additional cost in the filtering stage, provided the first stage remains memory-bound.

The authors emphasize that their theoretical analysis provides a rigorous foundation for parameter selection, offering a user-friendly interface (approx_top_k(array, K, recall_target)) that requires no manual tuning. While the empirical evaluation is focused on TPUv5e, the authors assert that the algorithm, performance modeling, and analysis are general and applicable to other hardware platforms like GPUs. The work demonstrates that principled algorithmic generalization, combined with hardware-aware implementation (fusion), can eliminate Top-K as a bottleneck in large-scale model training and inference.