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

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 top-$k$ algorithm with probabilistic correctness guarantees. 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 with probability at least $1-\epsilon$, where $\epsilon >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.5\times$-$15\times$ speedups over the highly optimized PyTorch top-$k$ 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, so the returned set can always be exact and $\epsilon$ bounds only the frequency of a second pass. We further demonstrate its impact on training BatchTopK Sparse Autoencoders (SAEs), where top-$k$ selection constitutes a significant portion of the training cost. The source code is provided in the supplementary materials and will be made publicly available.