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

Prof-K is a probabilistic one-pass filtering algorithm for efficient top-k selection, introduced by Tadeusz Dziarmaga, Witold Sikora, Łukasz Struski, Jacek Tabor, and Marcin Mazur from Jagiellonian University. The paper addresses the fundamental computational primitive of top-k selection, which has applications in databases, information retrieval, signal processing, and modern machine learning workloads including sparse activations and attention pruning.

The paper identifies that as data sizes grow, existing top-k 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. The challenge is most visible in the large-N, small-k regime, where even when only a tiny fraction of elements is needed, exact algorithms must still determine the precise global decision boundary across the full tensor. On modern GPUs, this cost is dominated primarily by memory movement and synchronization.

Prof-K proceeds in four stages:

1. Sampling: Uniformly sample S indices without replacement from [N]

2. Threshold Estimation: Sort the sampled values and define a threshold τ as the t-th largest sampled value

3. Filtering: Stream through the full input once, appending every element satisfying xi ≥ τ to a candidate buffer B of capacity M

4. Refinement: Run an exact top-k routine on the buffer B to recover the final answer, with fallback to exact top-k computation if fewer than k candidates are collected or if the buffer overflows

The central design goal is to choose parameters (S, t, M) so that the number of retained elements R satisfies k ≤ R ≤ M with probability at least 1 − ε, where ε is user-specified.

The paper derives distribution-agnostic guarantees based on the combinatorial properties of uniform sampling without replacement. Key theoretical results include:

Lemma 1 (Threshold rank distribution): The rank R of the sample-based threshold in the full population follows the negative hypergeometric distribution, with E[R] = (N+1)t/(S+1) and Var(R) = t(S−t+1)(N+1)(N−S)/((S+1)2(S+2)).

Lemma 2 (Normal approximation): For large N and S, R approximately follows N(Nq, N2q(1−q)f/S), where f = 1 − S/N is the finite-population correction.

Theorem 3 (Total failure probability bound): With threshold choice t according to Equation (6) and buffer size M = ⌈ck⌉ according to Equation (11), the total failure probability of Prof-K is at most ε, where ε = εA + εB is split between recall failures and buffer overflow.

Corollary 4 (Closed-form buffer size): M = k + z√(k(N−k)/S)√(1−S/N), where z = zA + zB.

Theorem 5 (Optimal sample size): When S ≪ N, the optimal sample size is S* = (z·τB·log(k)/(2τA))(2/3) · (k(N−k))(1/3), which simplifies to S* ∝ (kN)(1/3) in the common regime k ≪ N. For example, with N = 109 and k = 100, S* ≈ 4,600.

The paper emphasizes several important design choices:

• Distribution-agnostic guarantees: Because the analysis depends only on ranks induced by uniform sampling, the resulting guarantees are distribution-agnostic and remain valid under heavy-tailed or adversarial inputs.

• Fallback as safety mechanism: Rather than overprovisioning S or M in pathological regimes, Prof-K reverts to exact top-k whenever probabilistic conditions are not met economically, with amortized cost O(ε).

• Parameter selection: The algorithm restricts S to [Smin, Smax] where Smin = 32768 avoids high-variance threshold estimates and Smax = 131,072 keeps sampling overhead negligible. When αS < −ln(εA), the algorithm conservatively sets t = 1 (using sample maximum as threshold).

The paper evaluates Prof-K on synthetic benchmarks and real machine learning workloads:

Synthetic benchmarks: On DGX H100 and RTX 3060 GPUs, with tensor sizes N ∈ [222, 230] and selection sizes K ∈ [25, 219], Prof-K achieves:

• 1.5×–10× speedups over PyTorch topk and RadiK implementations

• Largest gains in the large-scale, small-to-moderate-k regime

• Performance invariant across Uniform, Standard Normal, and heavy-tailed Pareto distributions

• Negligible auxiliary memory consumption compared to radix-based methods requiring O(N) additional space

BatchTopK Sparse Autoencoder (SAE) training: Using activations from EleutherAI/pythia-70m-deduped at the MLP submodule of layer 1, training BatchTopK SAEs on OpenWebText with batch size B = 4096, learning rate 3×10−4, target sparsity k = 32, and dictionary size ddict = 12288:

• Prof-K reduces mean top-k time from approximately 2.56 ms to 1.13 ms (2.27× kernel-level speedup)

• Mean step time decreases from approximately 30.73 ms to 29.42 ms (1.044× end-to-end speedup)

• Total training-time reduction of approximately 4.25%

• Quality metrics (NMSE, cross-entropy degradation, fraction of variance explained, average active latents) remain nearly indistinguishable from exact top-k throughout training

The paper summarizes three main benefits of Prof-K:

1. Efficiency in the large-N, small-to-moderate-k regime: Reduces exact selection over N elements to exact refinement on a small compact candidate buffer

2. Distribution-agnostic guarantees: Depends only on ranks rather than values, making it robust to heavy-tailed or adversarial inputs

3. Flexible accuracy-speed tradeoffs: Users can tune sampling and buffer parameters to trade exact recovery for additional speed while retaining a fallback path to exact top-k

The paper acknowledges that Prof-K's benefits become less pronounced for smaller input sizes N, relies on fallback to exact top-k in rare cases, is currently limited to single-node GPU settings, and requires additional care when applied in batched versions since the probability of at least one failure increases with batch size.

Improvement 1: Probabilistic Top-k Kernel for Sparse Attention and MoE Routing

I can integrate Prof-K as a drop-in replacement for exact top-k in transformer architectures, specifically for sparse attention (e.g., selecting top-k keys per query) and Mixture-of-Experts (MoE) routing (e.g., selecting top-k experts per token). The improved AI system will:

• Reduce memory bandwidth and synchronization overhead during attention pruning, enabling longer context windows or larger batch sizes on the same GPU.

• Maintain near-identical model quality (e.g., perplexity, downstream accuracy) because Prof-K’s distribution-agnostic guarantees ensure the selected top-k set is correct with high probability, and fallback to exact top-k prevents catastrophic failures.

• Achieve 1.5–10× faster top-k computation in the large-N, small-k regime typical of attention (N = sequence length × head dimension, k = 8–128), directly reducing per-layer latency.

Improvement 2: Adaptive Sampling for Dynamic Data Distributions in Online Learning

I can use Prof-K’s rank-based threshold estimation to build an adaptive sampler for online learning systems (e.g., recommender systems, active learning). The improved AI system will:

• Dynamically adjust the sample size S and threshold rank t based on observed data stream statistics, using the closed-form buffer size M = k + z√(k(N−k)/S)√(1−S/N) to bound failure probability without retraining.

• Handle heavy-tailed or adversarial input distributions (e.g., click-through rates, sensor anomalies) without degradation, since guarantees depend only on ranks, not values.

• Provide a tunable accuracy-speed knob: users can increase ε to trade exactness for higher throughput, with the fallback path ensuring no catastrophic loss of critical top-k elements.

Improvement 3: Memory-Efficient Top-k for Large-Scale Embedding Retrieval

I can apply Prof-K to embedding-based retrieval systems (e.g., nearest neighbor search in vector databases). The improved AI system will:

• Replace exact sorting over millions of embeddings with a single-pass filtering step, reducing auxiliary memory from O(N) (as in radix-based methods) to O(M) where M ≈ k + z√(k(N−k)/S). This enables serving billion-scale embedding tables on memory-constrained GPUs.

• Maintain high recall (≥1−ε) for top-k retrieval, with the fallback to exact top-k triggered only when the sample-based threshold is unreliable (e.g., when k is large relative to N).

• Support batched queries by allocating a shared buffer per batch, with a batch-level failure probability bound (ε batch = 1 − (1−ε) B) to guide buffer sizing.

Improvement 4: Hardware-Aware Auto-Tuning for Top-k in Deep Learning Compilers

I can use Prof-K’s theoretical results (optimal sample size S* ∝ (kN)(1/3), buffer size M) to build an auto-tuning module in deep learning compilers (e.g., XLA, TVM). The improved AI system will:

• Automatically select S, t, and M based on tensor shape, GPU memory bandwidth, and target failure probability ε, eliminating manual hyperparameter search.

• Generate fused kernels that combine sampling, threshold estimation, filtering, and refinement into a single CUDA kernel, reducing kernel launch overhead and synchronization points.

• Profile at runtime to adjust parameters if the observed failure rate exceeds ε (e.g., due to unexpected data skew), ensuring robust performance across diverse workloads.

Improvement 5: Probabilistic Top-k for Distributed and Federated Learning

I can extend Prof-K to distributed settings where each worker holds a shard of the data. The improved AI system will:

• Have each worker run Prof-K locally on its shard, producing a candidate buffer of size M, then merge buffers across workers using a hierarchical exact top-k. This reduces communication cost from O(N) to O(M × workers).

• Use the rank-based guarantees to bound the global failure probability, even when shard sizes are unequal or data distributions differ across workers (e.g., non-IID federated learning).

• Provide a fallback to full exact top-k only when the merged buffer underflows, ensuring correctness in adversarial sharding scenarios.

Improvement 6: Real-Time Anomaly Detection with Guaranteed Top-k Alerts

I can use Prof-K to build a streaming anomaly detection system that must identify the top-k most anomalous events (e.g., network intrusions, system failures) from a high-velocity data stream. The improved AI system will:

• Process each new batch of N events in O(N + M log M) time instead of O(N log N), enabling real-time alerting on 106–109 events per second.

• Provide a probabilistic guarantee that the true top-k anomalies are included with probability ≥1−ε, with the fallback path ensuring no missed alerts when the sample-based threshold is too high.

• Tune ε dynamically based on the cost of false negatives (e.g., higher ε during low-traffic periods, lower ε during critical windows) without recomputing the full top-k.

Improvement 7: Accelerated Training of Sparse Models via BatchTopK

I can replace the exact BatchTopK operation in sparse autoencoders (SAEs) and sparse transformers with Prof-K. The improved AI system will:

• Reduce per-step training time by 4.25% end-to-end (as demonstrated in the paper) by cutting top-k kernel time from 2.56 ms to 1.13 ms on a batch of 4096, enabling larger models or longer training runs within the same compute budget.

• Maintain identical model quality metrics (NMSE, cross-entropy, variance explained) because Prof-K’s fallback ensures exactness in rare failure cases, and the probabilistic errors are negligible during training.

• Scale to larger dictionary sizes (e.g., d dict = 65536) where exact top-k becomes the bottleneck, by using the optimal sample size S* ≈ (kN)(1/3) to keep sampling overhead low.

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.