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

2026-08-12

AI Summary

Existing Top-k algorithms suffer from either low efficiency or poor robustness on large-scale data: exact methods incur high computational overhead, while approximate approaches exhibit unstable performance under adversarial or heavy-tailed distributions. This work proposes Prof-K, the first distribution-agnostic, single-pass streaming Top-k algorithm. Prof-K adaptively estimates a threshold via random sampling to filter input into a compact buffer, on which an exact Top-k computation is performed, provably recovering the true Top-k result with probability at least $1 - \varepsilon$. The method allows users to specify an error bound and offers a tunable trade-off between accuracy and speed. Experiments demonstrate that Prof-K achieves 1.5–10× speedups over PyTorch’s topk and RadiK for moderate values of $k$, maintains robustness against adversarial inputs, and has been successfully deployed in BatchTopK sparse autoencoder training, substantially reducing computational costs.

Abstract

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 - $ε$, where $ε$ > 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-10x 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. By relaxing the recall target (e.g., recovering 95% of the true top-k values), Prof-K additionally provides a principled accuracy-speed trade-off. We further demonstrate its impact on training BatchTopK Sparse Autoencoders (SAEs), where top-k selection constitutes a significant portion of the training cost.

Problem

• top-k selection
• scalability
• adversarial inputs
• heavy-tailed distributions
• computational efficiency

Innovation

• one-pass filtering
• probabilistic top-k
• distribution-agnostic
• adaptive thresholding
• sparse autoencoders