# Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection

arXiv: https://arxiv.org/abs/2608.12573 · HTML: https://mlsanity.com/paper/2608.12573/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 33 references that are in the map.
- First posted on arXiv 12 Aug 2026.

## Builds on (most cited first, 16 of 16)

- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [Learning Multi-Level Features with Matryoshka Sparse Autoencoders](https://mlsanity.com/paper/2503.17547/) (2025)
- [Understanding Top-k Sparsification in Distributed Deep Learning](https://mlsanity.com/paper/1911.08772/) (2019)
- [Overcoming Sparsity Artifacts in Crosscoders to Interpret Chat-Tuning](https://mlsanity.com/paper/2504.02922/) (2025)
- [An Efficient Statistical-based Gradient Compression Technique for Distributed Training Systems](https://mlsanity.com/paper/2101.10761/) (2021)
- [On the Utility of Gradient Compression in Distributed Training Systems](https://mlsanity.com/paper/2103.00543/) (2021)
- [Near-Optimal Sparse Allreduce for Distributed Deep Learning](https://mlsanity.com/paper/2201.07598/) (2022)
- [TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s](https://mlsanity.com/paper/2206.14286/) (2022)
- [RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs](https://mlsanity.com/paper/2409.00822/) (2024)
- [Dr. Top-k: Delegate-Centric Top-k on GPUs](https://mlsanity.com/paper/2109.08219/) (2021)
- [RadiK: Scalable and Optimized GPU-Parallel Radix Top-K Selection](https://mlsanity.com/paper/2501.14336/) (2025)
- [Empirical Analysis on Top-k Gradient Sparsification for Distributed Deep Learning in a Supercomputing Environment](https://mlsanity.com/paper/2209.08497/) (2022)

## Abstract (arXiv)

Top-k selection is a fundamental computational primitive with applications spanning databases, information retrieval, signal processing, and modern machine learning workloads, including sparse activations and attention pruning. As data sizes grow, existing approaches become inefficient: exact methods incur high memory and compute overhead, while approximate methods often rely on brittle heuristics that degrade under adversarial or heavy-tailed inputs. In this paper, we introduce Prof-K, a fast, scalable, and distribution-agnostic algorithm for exact top-k selection. Prof-K performs a single-pass filtering procedure: a small random sample estimates an adaptive threshold, the N input elements are streamed once into a compact buffer, and an exact top-k routine on this buffer recovers the true top-k elements on the first attempt with probability at least 1-varepsilon, where varepsilon>0 is user specified. We derive high-probability guarantees for correctness and buffer size, together with an approximately optimal sample size that minimizes overhead as a function of N and k. Empirically, Prof-K achieves 1.5x-15x speedups over the highly optimized PyTorch topk and recent RadiK implementations, with the largest gains in the large-scale, small-to-moderate-k regime where prior methods struggle most. Unlike previous approaches, these guarantees hold independently of the input distribution, ensuring robustness to adversarial settings. A run-time check detects the rare failures and triggers a retry, so the returned set is always exact and varepsilon bounds only the probability of requiring an additional pass. We further demonstrate its impact on training BatchTopK Sparse Autoencoders (SAEs), where top-k selection constitutes a significant portion of the training cost.