摘要：Top-k 选择是一项基础计算原语，应用范围涵盖数据库、信息检索、信号处理以及现代机器学习工作负载，包括稀疏激活和注意力剪枝。随着数据规模增长，现有方法变得低效：精确方法会带来较高的内存和计算开销，而近似方法往往依赖脆弱的启发式策略，在对抗性输入或重尾输入下性能会下降。本文提出 Prof-K，这是一种快速、可扩展且与数据分布无关的 Top-k 算法，并提供概率正确性保证。Prof-K 执行单遍过滤过程：先利用一个小型随机样本估计自适应阈值，将 N 个输入元素仅流式遍历一次并写入紧凑缓冲区，随后在该缓冲区上运行精确的 Top-k 过程，以至少 $1 - \epsilon$ 的概率恢复真实的 Top-k 元素，其中 $\epsilon > 0$ 由用户指定。我们推导了关于正确性和缓冲区大小的高概率保证，并给出了一个随 N 和 k 变化、可将额外开销最小化的近似最优样本规模。实验表明，与经过高度优化的 PyTorch topk 以及近期的 RadiK 实现相比，Prof-K 可实现 1.5 倍至 10 倍的加速；在大规模、k 较小至中等的场景中提升最为显著，而这正是先前方法最难处理的领域。与以往方法不同，这些保证独立于输入分布成立，从而确保了其在对抗性环境下的鲁棒性。通过放宽召回率目标，例如仅恢复真实 Top-k 值中的 95%，Prof-K 还提供了有原则的精度—速度权衡。我们进一步展示了它对 BatchTopK Sparse Autoencoder（SAE）训练的影响；在该任务中，Top-k 选择构成了训练成本的重要部分。

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 - $\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.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.