Prof-K：用于高效 Top-k 选择的概率单遍过滤

arXiv:2608.12573v2 公告类型：替换 摘要：Top-k 选择是一种基础计算原语，应用于数据库、信息检索、信号处理以及现代机器学习工作负载，包括稀疏激活和注意力剪枝。随着数据规模不断增长，现有方法的效率逐渐降低：精确方法会产生较高的内存和计算开销，而近似方法通常依赖脆弱的启发式策略，在对抗性输入或重尾输入下性能会下降。在本文中，我们提出了 Prof-K，一种快速、可扩展且与数据分布无关的精确 Top-k 选择算法。Prof-K 执行单遍过滤过程：通过一个小型随机样本估计自适应阈值，将 N 个输入元素以流式方式遍历一次并存入紧凑缓冲区，然后在该缓冲区上运行精确的 Top-k 例程，以至少 $1-\varepsilon$ 的概率在首次尝试时恢复真实的 Top-k 元素，其中 $\varepsilon>0$ 由用户指定。我们推导了正确性和缓冲区大小的高概率保证，并给出了近似最优的样本大小，使开销能够根据 N 和 k 实现最小化。在实证评估中，与高度优化的 PyTorch topk 和近期的 RadiK 实现相比，Prof-K 实现了 1.5 倍至 15 倍的加速；其中，在先前方法最难应对的大规模、小到中等 k 的场景下，性能提升最大。与以往方法不同，这些保证与输入分布无关，因此能够确保在对抗性环境下的稳健性。运行时检查会检测罕见的失败情况并触发重试，因此返回的集合始终是精确的，而 $\varepsilon$ 仅限定需要额外遍历的概率。我们还展示了 Prof-K 对 BatchTopK Sparse Autoencoders（SAEs）训练的影响；在该训练过程中，Top-k 选择占据了相当大比例的训练成本。