arXiv：cs.LG（机器学习，全量分类）· Tadeusz Dziarmaga, Witold Sikora, {\L}ukasz Struski, Jacek Tabor, Marcin Mazur·· 15 小时前AI 评分44
Prof-K：面向高效 Top-k 选择的概率单遍过滤算法
Prof-K: Probabilistic One-Pass Filtering for Efficient Top-k Selection
AI 导读
Prof-K 是一种分布无关的精确 Top-k 选择算法，通过单遍过滤实现：先用小随机样本估计自适应阈值，将 N 个元素流式写入紧凑缓冲区，再对缓冲区做精确 Top-k，首次尝试即以至少 1-ε 的概率恢复真实 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 algorithm for exact top-k selection. 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 on the first attempt with probability at least $1-\varepsilon$, where $\varepsilon>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-15x 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. A run-time check detects the rare failures and triggers a retry, so the returned set is always exact and $\varepsilon$ bounds only the probability of requiring an additional pass. We further demonstrate its impact on training BatchTopK Sparse Autoencoders (SAEs), where top-k selection constitutes a significant portion of the training cost.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2608.12573 [cs.LG] |
| (or arXiv:2608.12573v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2608.12573 arXiv-issued DOI via DataCite |
Submission history
From: Tadeusz Dziarmaga [view email]
[v1]
Wed, 12 Aug 2026 20:31:33 UTC (4,607 KB)
[v2]
Wed, 30 Sep 2026 19:05:02 UTC (320 KB)
来源：arXiv：cs.LG（机器学习，全量分类） · arxiv.org