一种快速可靠的单遍筛选算法,高效选出前k大值。
Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection

- 用随机采样自适应设定阈值,单次遍历数据流并压缩缓冲区。
- 在用户指定误差率ε下,以至少1-ε概率准确找回真实top-k元素。
- 适用于任意数据分布,特别适合大规模小k场景,适合模型训练优化。
Top-k选择是数据库、信息检索、信号处理及现代机器学习任务(如稀疏激活与注意力剪枝)中的基础计算操作。随着数据规模增长,现有方法效率下降:精确方法开销大,近似方法常依赖脆弱启发式,在对抗性或重尾输入下性能退化。本文提出Prof-K,一种快速、可扩展且分布无关的top-k算法,具备概率正确性保障。Prof-K执行一次遍历过滤:小规模随机采样估计自适应阈值,将N个输入元素流式处理至紧凑缓冲区,再在该缓冲区上使用精确top-k算法,以不低于1-ε的概率恢复真实top-k元素,ε>0由用户指定。我们推导了正确性与缓冲区大小的高概率保证,并给出近似最优样本量,以最小化关于N和k的开销。实验表明,Prof-K相比高度优化的PyTorch topk和近期RadiK实现,提速1.5倍至10倍,尤其在大规模、小至中等k的场景下优势显著。该方法的保证独立于输入分布,对对抗环境具有鲁棒性。通过放宽召回目标(如恢复95%真实top-k),Prof-K还提供严谨的精度-速度权衡。我们进一步展示了其在BatchTopK稀疏自编码器(SAEs)训练中的影响,其中top-k选择占训练成本较大比例。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。