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

# Probabilistic One-Pass Filtering for Efficient Top- k Selection
Source: [https://arxiv.org/html/2608.12573](https://arxiv.org/html/2608.12573)
Tadeusz DziarmagaWitold Sikora11footnotemark:1Affiliation:Jagiellonian UniversityŁukasz StruskiAffiliation:Jagiellonian UniversityJacek TaborAffiliation:Jagiellonian UniversityMarcin MazurAffiliation:Jagiellonian University

###### Abstract

Top\-kkselection 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\-kkalgorithm with probabilistic correctness guarantees\. Prof\-K performs a single\-pass filtering procedure: a small random sample estimates an adaptive threshold, theNNinput elements are streamed once into a compact buffer, and an exact top\-kkroutine on this buffer recovers the true top\-kkelements with probability at least1−ε1\-\\varepsilon, whereε\>0\\varepsilon\>0is 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 ofNNandkk\. Empirically, Prof\-K achieves1\.5×1\.5\\times–10×10\\timesspeedups over the highly optimized PyTorchtopkand recent RadiK implementations, with the largest gains in the large\-scale, small\-to\-moderate\-kkregime 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\-kkvalues\), Prof\-K additionally provides a principled accuracy–speed trade\-off\. We further demonstrate its impact on training BatchTopK Sparse Autoencoders \(SAEs\), where top\-kkselection constitutes a significant portion of the training cost\.

Figure 1:The core idea of our Prof\-K algorithm\.Top:Full data of sizeNNsorted in descending order; the true top\-kkboundary is marked by a vertical dashed line\.Middle:A random sample \(solid outlines\) is used to compute a thresholdτ\\tau\(green dashed line\) as a conservative estimate of the\(k/N\)\(k/N\)\-quantile\.Bottom:The threshold filter is applied in parallel to the full data\. Elements≥τ\\geq\\tau\(green\) are retained in a compact candidate buffer, while elements belowτ\\tauare discarded\. This approach ensures all true top\-kkelements are preserved with high probability while drastically reducing the workload for the final selection stage\.## 1Introduction

Top\-kkselection is a common systems operation that arises whenever only a small subset of the largest elements must be retained from a much larger collection\[[1](https://arxiv.org/html/2608.12573#bib.bib11),[12](https://arxiv.org/html/2608.12573#bib.bib8)\]\. As tensor sizes continue to grow, the cost of repeated top\-kkexecution becomes increasingly significant in large\-scale GPU workloads, particularly in sparse computation settings where only a small fraction of values is ultimately retained\[[13](https://arxiv.org/html/2608.12573#bib.bib5),[8](https://arxiv.org/html/2608.12573#bib.bib6),[5](https://arxiv.org/html/2608.12573#bib.bib7),[3](https://arxiv.org/html/2608.12573#bib.bib1)\]\.

The challenge is most visible in the large\-NN, small\-kkregime\. 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, making top\-kkdisproportionately expensive relative to the surrounding computation\[[2](https://arxiv.org/html/2608.12573#bib.bib9),[6](https://arxiv.org/html/2608.12573#bib.bib2),[16](https://arxiv.org/html/2608.12573#bib.bib13),[15](https://arxiv.org/html/2608.12573#bib.bib3),[9](https://arxiv.org/html/2608.12573#bib.bib4)\]\.

In this work, we introduce Prof\-K, a probabilistic one\-pass filtering algorithm for efficient top\-kkselection\. The method first draws a small uniform sample to estimate a conservative threshold, then streams once through the input tensor while retaining only elements above this threshold in a compact candidate buffer\. Exact top\-kkrefinement is finally applied only to the retained candidates\. 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\. A schematic summary of Prof\-K is presented in[Figure1](https://arxiv.org/html/2608.12573#S0.F1)\.

We derive explicit parameter choices controlling both recall failure and buffer overflow, together with an approximately optimal sample size scaling as\(k​N\)1/3\(kN\)^\{1/3\}in the sparse regime\. Empirically, Prof\-K achieves substantial speedups over optimized baselines such as PyTorchtopkand RadiK\[[9](https://arxiv.org/html/2608.12573#bib.bib4)\], particularly for very large tensors and small retained fractions\.

We additionally evaluate Prof\-K in BatchTopK Sparse Autoencoder \(SAE\)\[[3](https://arxiv.org/html/2608.12573#bib.bib1)\]training\. Even on a relatively small language model, we observe measurable reductions in end\-to\-end training time without changes in reconstruction quality or sparsity behavior\. Since the flattened BatchTopK selection problem scales directly with model and dictionary size, we expect the advantages of Prof\-K to become even stronger in larger SAE settings such as GPT\-2 Small\[[11](https://arxiv.org/html/2608.12573#bib.bib14)\]and Gemma 2 2B\[[14](https://arxiv.org/html/2608.12573#bib.bib16)\]\.

Overall, Prof\-K provides three main benefits:

- •Efficiency in the large\-NN, small\-to\-moderate\-kkregime\.Prof\-K reduces exact selection overNNelements to exact refinement on a small compact candidate buffer, with benefits that grow asNNincreases\.
- •Distribution\-agnostic guarantees\.Unlike quantile\-estimation methods that assume well\-behaved distributions, Prof\-K derives distribution\-agnostic guarantees from random sampling and depends only on ranks rather than values, making it robust to heavy\-tailed or adversarial inputs\.
- •Flexible accuracy\-speed tradeoffs\.Users can tune the sampling and buffer parameters of Prof\-K to trade exact recovery for additional speed while retaining a fallback path to exact top\-kkwhen needed\.

## 2Related Work

Top\-kkselection is a fundamental primitive with applications in databases, information retrieval, signal processing, and modern machine learning systems\. Classical exact methods are based on sorting, heaps, partition\-based selection, or partial sorting routines such as quickselect andnth\_element\. Efficient top\-kkquery execution has long been studied in database systems and parallel hardware settings\[[1](https://arxiv.org/html/2608.12573#bib.bib11),[12](https://arxiv.org/html/2608.12573#bib.bib8)\]\. In large\-scale GPU workloads, however, the dominant cost is often memory movement rather than comparison complexity alone, making exact top\-kkexpensive even in optimized libraries such as PyTorch\.

A substantial line of work studies efficient exact top\-kkselection on GPUs\. Early work by Alabi et al\.\[[2](https://arxiv.org/html/2608.12573#bib.bib9)\]developed GPU\-specifickk\-selection algorithms optimized for massively parallel architectures\. Dr\. Top\-k\[[6](https://arxiv.org/html/2608.12573#bib.bib2)\]introduces a delegate\-centric design that improves load balancing and reduces redundant work among GPU threads\. Zhang et al\.\[[16](https://arxiv.org/html/2608.12573#bib.bib13)\]provide a comprehensive study of parallel GPU top\-kkalgorithms and propose optimized implementations for different operating regimes\. RTop\-K\[[15](https://arxiv.org/html/2608.12573#bib.bib3)\]focuses on efficient row\-wise top\-kkselection and demonstrates strong performance for moderate\-size neural network workloads\. Radix\-based approaches such as RadiK\[[9](https://arxiv.org/html/2608.12573#bib.bib4)\]further exploit floating\-point representations to accelerate exact selection\. In contrast, Prof\-K reduces the size of the exact selection problem itself through probabilistic filtering, making it particularly effective for very largeNNand relatively smallkk\.

Approximate and probabilistic selection methods reduce computational cost using sampling, sketching, or approximate quantile estimation, often at the cost of weaker guarantees or assumptions on the input distribution\. Quantile summary methods such as the Greenwald–Khanna\[[7](https://arxiv.org/html/2608.12573#bib.bib10)\]algorithm provide efficient streaming approximations with bounded error guarantees\. Prof\-K differs in that its guarantees depend only on the combinatorial properties of uniform sampling without replacement\. The analysis therefore depends on ranks rather than values, yielding distribution\-agnostic probabilistic guarantees\.

Top\-kkselection is increasingly important in sparse machine learning workloads, including mixture\-of\-experts routing, sparse attention, activation pruning, and sparse autoencoders\[[13](https://arxiv.org/html/2608.12573#bib.bib5),[8](https://arxiv.org/html/2608.12573#bib.bib6),[5](https://arxiv.org/html/2608.12573#bib.bib7),[3](https://arxiv.org/html/2608.12573#bib.bib1)\]\. These workloads commonly operate in the regime targeted by Prof\-K: extremely large tensors with small retained fractions\. Rather than replacing exact selection entirely, Prof\-K uses sampling to construct a conservative threshold and then applies exact top\-kkrefinement on a compact candidate buffer, making it complementary to optimized GPU kernels such as Dr\. Top\-k, RTop\-K, and RadiK\.

## 3Problem Setup and Algorithm Overview

Given an input vectorx∈ℝNx\\in\\mathbb\{R\}^\{N\}and an integerk≥1k\\geq 1, our goal is to return thekklargest entries ofxxtogether with their indices\. Letα=k/N\\alpha=k/Ndenote the target fraction of retained elements\. For batched inputs with batch sizeBB, the procedure is applied in parallel to each batch element\.

### The Prof\-K Algorithm

Prof\-K proceeds in four stages, designed to minimize memory traffic while providing probabilistic correctness guarantees:

1. 1\.Sampling\.Uniformly sampleSSindices without replacement from\[N\]=\{1,…,N\}\[N\]=\\\{1,\\dots,N\\\}\. Let𝒮⊂\[N\]\\mathcal\{S\}\\subset\[N\]denote the sampled indices\.
2. 2\.Threshold Estimation\.Sort the sampled valuesx\|𝒮x\|\_\{\\mathcal\{S\}\}in descending order\. For a chosen rankt∈\[S\]t\\in\[S\], define the thresholdτ=st\\tau=s\_\{t\}, wherests\_\{t\}is thett\-th largest sampled value\.
3. 3\.Filtering\.Stream through the full input once, appending every element satisfyingxi≥τx\_\{i\}\\geq\\tauto a candidate bufferℬ\\mathcal\{B\}of capacityMM, which is a small constant multiple ofkk\.
4. 4\.Refinement\.Run an exact top\-kkroutine on the bufferℬ\\mathcal\{B\}to recover the final answer\. If fewer thankkcandidates are collected, or if the buffer overflows before the scan completes, Prof\-K falls back to an exact top\-kkcomputation on the full input\.

The central design goal is to choose the parameters\(S,t,M\)\(S,t,M\)so that the number of retained elementsR=\|\{i∈\[N\]:xi≥τ\}\|R=\|\\\{i\\in\[N\]:x\_\{i\}\\geq\\tau\\\}\|satisfiesk≤R≤Mk\\leq R\\leq Mwith probability at least1−ε1\-\\varepsilon, where0<ε≪10<\\varepsilon\\ll 1is user specified\. When this holds, the final exact top\-kkoperates on only a small multiple ofkkelements rather than allNNentries\.

### Key Insight: Why Sampling Works

The core idea is that a random sample provides an unbiased estimate of quantile positions\. If we select the\(t/S\)\(t/S\)\-quantile from the sample, it approximates the\(t/S\)\(t/S\)\-quantile of the full population\. By choosingttslightly larger thanS​α=S​k/NS\\alpha=Sk/N, we obtain a threshold that is conservatively low, ensuring all true top\-kkelements pass the filter with high probability, while still excluding most of the remaining elements\.

## 4Theoretical Analysis

We now formalize the probabilistic guarantees underlying the Prof\-K algorithm\. Our analysis proceeds in three steps: \(1\) characterizing the distribution of the estimated threshold rank, \(2\) deriving parameter choices that bound the failure probability, and \(3\) optimizing the sample size to minimize computational overhead\. The omitted proofs are provided in Appendix[C](https://arxiv.org/html/2608.12573#A3)\.

### Distribution of the Threshold Rank

The first key observation is that the rank of our sample\-based threshold in the full population follows a well\-understood distribution, regardless of the input values themselves\.

###### Lemma 1\(Threshold rank distribution\)\.

Let𝒮⊂\[N\]\\mathcal\{S\}\\subset\[N\]be a uniform sample of sizeSSdrawn without replacement\. Letτ\\taube thett\-th largest sampled value, and letRRdenote its rank in the fully sorted input \(rank11is the largest\)\. ThenRRfollows the negative hypergemoetric distribution, i\.e\., forr∈\[t,N−S\+t\]r\\in\[t,N\-S\+t\],

Pr⁡\(R=r\)=\(r−1t−1\)​\(N−rS−t\)\(NS\)\.\\Pr\(R=r\)=\\frac\{\\binom\{r\-1\}\{t\-1\}\\binom\{N\-r\}\{S\-t\}\}\{\\binom\{N\}\{S\}\}\.\(1\)Moreover,

𝔼⁡\[R\]=N\+1S\+1​t,Var⁡\(R\)=t⁡\(S−t\+1\)​\(N\+1\)​\(N−S\)\(S\+1\)2​\(S\+2\)\.\\mathbb\{E\}\[R\]=\\frac\{N\+1\}\{S\+1\}\\,t,\\;\\;\\;\\mathrm\{Var\}\(R\)=\\frac\{t\(S\-t\+1\)\(N\+1\)\(N\-S\)\}\{\(S\+1\)^\{2\}\(S\+2\)\}\.\(2\)

This result is powerful because*it depends only on the ranks, not on the actual values inxx*\. Whether the input is uniformly distributed, heavy\-tailed, or adversarially constructed, the same probabilistic bounds apply\. For practical parameter regimes, a Gaussian approximation provides convenient closed\-form expressions\.

###### Lemma 2\(Normal approximation\)\.

SupposeN,S→∞N,S\\to\\inftywithS/N→f∞∈\(0,1\)S/N\\to f\_\{\\infty\}\\in\(0,1\)\. Letq=t/Sq=t/S\. Then

R−N​qN​q⁡\(1−q\)S​\(1−SN\)→𝑑𝒩⁡\(0,1\)\.\\frac\{R\-Nq\}\{N\\sqrt\{\\frac\{q\(1\-q\)\}\{S\}\\left\(1\-\\frac\{S\}\{N\}\\right\)\}\}\\xrightarrow\{d\}\\mathcal\{N\}\(0,1\)\.\(3\)Hence, for large finiteN,SN,S,

R≈𝒩⁡\(N​q,N2​q⁡\(1−q\)S​f\),f=1−SN\.R\\approx\\mathcal\{N\}\\\!\\left\(Nq,\\;N^\{2\}\\frac\{q\(1\-q\)\}\{S\}f\\right\),\\quad f=1\-\\frac\{S\}\{N\}\.\(4\)

The factorf=1−S/Nf=1\-S/Nis the finite\-population correction: sampling without replacement reduces variance when the sample constitutes a non\-negligible fraction of the population\.

### Failure Modes and Budget Allocation

Prof\-K can fail in two distinct ways, each requiring separate treatment:

1. 1\.Insufficient Recall\.The threshold is too high \(R<kR<k\), so fewer thankkvalid candidates survive filtering\.
2. 2\.Buffer Overflow\.Too many elements exceed the threshold \(R\>MR\>M\), and some true top\-kkentries fails to enter the finite buffer\.

We allocate a total failure budgetε=εA\+εB\\varepsilon=\\varepsilon\_\{A\}\+\\varepsilon\_\{B\}across these events, withεA\\varepsilon\_\{A\}controlling recall failures andεB\\varepsilon\_\{B\}controlling overflow\.

### Controlling Insufficient Recall

This failure mode corresponds to the event\{R<k\}\\\{R<k\\\}\. Using the Gaussian approximation in[Equation4](https://arxiv.org/html/2608.12573#S4.E4), we obtain

Pr⁡\(R<k\)≈Φ⁡\(k−N​qN​q⁡\(1−q\)S​f\)≤εA,\\Pr\(R<k\)\\approx\\Phi\\\!\\left\(\\frac\{k\-Nq\}\{N\\sqrt\{\\frac\{q\(1\-q\)\}\{S\}f\}\}\\right\)\\leq\\varepsilon\_\{A\},\(5\)whereΦ\\Phiis the standard normal CDF\. We express the sampling quantile rank asq=α\+δq=\\alpha\+\\deltawithα=k/N\\alpha=k/N, and note that for smallδ\\deltawe havek−N​q=−N​δk\-Nq=\-N\\deltaandq⁡\(1−q\)≈α⁡\(1−α\)q\(1\-q\)\\approx\\alpha\(1\-\\alpha\)\. Substituting into the tail bound yieldsΦ\(−δ/α⁡\(1−α\)S​f\)≤εA\\Phi\\\!\\left\(\-\\delta/\\sqrt\{\\frac\{\\alpha\(1\-\\alpha\)\}\{S\}f\}\\right\)\\leq\\varepsilon\_\{A\}, which impliesδ≥zA​α⁡\(1−α\)​fS\\delta\\geq z\_\{A\}\\sqrt\{\\frac\{\\alpha\(1\-\\alpha\)f\}\{S\}\}withzA=Φ−1​\(1−εA\)z\_\{A\}=\\Phi^\{\-1\}\(1\-\\varepsilon\_\{A\}\)\. This leads to the threshold choice

q=α\+zA​α⁡\(1−α\)​fS,t=⌈S​α\+zA​S​α​\(1−α\)​f⌉,q=\\alpha\+z\_\{A\}\\sqrt\{\\frac\{\\alpha\(1\-\\alpha\)f\}\{S\}\},\\qquad t=\\left\\lceil S\\alpha\+z\_\{A\}\\sqrt\{S\\alpha\(1\-\\alpha\)f\}\\,\\right\\rceil,\(6\)which can be interpreted as shifting the expected sample rankS​αS\\alphaupward byzAz\_\{A\}standard deviations of its sampling distribution to reduce the probability of discarding true top\-kkelements; whenS≪NS\\ll N, this simplifies to the standard binomial formt≈α​S\+zA​α⁡\(1−α\)​St\\approx\\alpha S\+z\_\{A\}\\sqrt\{\\alpha\(1\-\\alpha\)S\}\. In the extreme sparse regimeα​S≪1\\alpha S\\ll 1, the Gaussian approximation is unreliable, and a conservative alternative is to sett=1t=1, yielding failure probabilityPr⁡\(Bin⁡\(S,α\)=0\)=\(1−α\)S≈e−α​S\\Pr\(\\Bin\(S,\\alpha\)=0\)=\(1\-\\alpha\)^\{S\}\\approx e^\{\-\\alpha S\}, so thatα​S≥−ln⁡εA\\alpha S\\geq\-\\ln\\varepsilon\_\{A\}ensures the desired bound; otherwise \(rarely\), the algorithm falls back to exact top\-kkcomputation\.

### Controlling Buffer Overflow

This failure mode is \(more than\) covered by the event\{R\>M\}\\\{R\>M\\\}\. LetM=c​kM=ckwith buffer multiplierc≥1c\\geq 1\. Using the threshold choice from[Equation6](https://arxiv.org/html/2608.12573#S4.E6), the expected number of retained elements is𝔼⁡\[R\]=N​q\\mathbb\{E\}\[R\]=Nq, which can be written as

𝔼⁡\[R\]=k\+N​zA​α⁡\(1−α\)​fS=k⁡\(1\+zAS​1−αα​f\),\\mathbb\{E\}\[R\]=k\+Nz\_\{A\}\\sqrt\{\\frac\{\\alpha\(1\-\\alpha\)f\}\{S\}\}=k\\\!\\left\(1\+\\frac\{z\_\{A\}\}\{\\sqrt\{S\}\}\\sqrt\{\\frac\{1\-\\alpha\}\{\\alpha\}f\}\\right\),\(7\)motivating the definition of the nominal buffer multiplier

c0=1\+zAS​1−αα​f\.c\_\{0\}=1\+\\frac\{z\_\{A\}\}\{\\sqrt\{S\}\}\\sqrt\{\\frac\{1\-\\alpha\}\{\\alpha\}f\}\.\(8\)The variability of the retained set is governed by

Var⁡\(R\)=σR=N​q⁡\(1−q\)S​f≈k​1−αα​S​f,\\Var\(R\)=\\sigma\_\{R\}=N\\sqrt\{\\frac\{q\(1\-q\)\}\{S\}f\}\\approx k\\sqrt\{\\frac\{1\-\\alpha\}\{\\alpha S\}f\},\(9\)so that by a Gaussian approximation the overflow probability satisfies

Pr⁡\(R\>c​k\)≈Φ⁡\(−\(c−c0\)​kσR\)=Φ⁡\(−\(c−c0\)​α​S\(1−α\)​f\)\.\\Pr\(R\>ck\)\\approx\\Phi\\\!\\left\(\-\\frac\{\(c\-c\_\{0\}\)k\}\{\\sigma\_\{R\}\}\\right\)=\\Phi\\\!\\left\(\-\(c\-c\_\{0\}\)\\sqrt\{\\frac\{\\alpha S\}\{\(1\-\\alpha\)f\}\}\\right\)\.\(10\)Enforcing this to be at mostεB\\varepsilon\_\{B\}yields

c=c0\+zB​\(1−α\)​fα​S,zB=Φ−1​\(1−εB\)\.c=c\_\{0\}\+z\_\{B\}\\sqrt\{\\frac\{\(1\-\\alpha\)f\}\{\\alpha S\}\},\\qquad z\_\{B\}=\\Phi^\{\-1\}\(1\-\\varepsilon\_\{B\}\)\.\(11\)

### Main Correctness Guarantee

Our main theoretical result follows immediately from the preceding analysis\.

###### Theorem 3\(Total failure probability bound\)\.

Assume a target failure budgetε=εA\+εB\\varepsilon=\\varepsilon\_\{A\}\+\\varepsilon\_\{B\}withεA,εB≪1\\varepsilon\_\{A\},\\varepsilon\_\{B\}\\ll 1\(typicallyεA=εB=ε/2\\varepsilon\_\{A\}=\\varepsilon\_\{B\}=\\varepsilon/2\)\. Choosettaccording to[Equation6](https://arxiv.org/html/2608.12573#S4.E6), and setM=⌈c​k⌉M=\\lceil ck\\rceil, whereccis given by[Equation11](https://arxiv.org/html/2608.12573#S4.E11)\. Then the total failure probability of Prof\-K is at mostε\\varepsilon\.

This result formalizes a practical design principle: one part of the failure budget controls threshold quality \(ensuring we do not miss true top\-kkelements\), and the other controls buffer adequacy \(ensuring we have room for all candidates\)\.

### Retrieving the Buffer Size

Combining the expressions forttandccyields a compact closed\-form for the buffer size\.

###### Corollary 4\(Closed\-form buffer size\)\.

Letz=zA\+zBz=z\_\{A\}\+z\_\{B\}\. Then

M=k\+z​k⁡\(N−k\)S​\(1−SN\)\.M=k\+z\\sqrt\{\\frac\{k\(N\-k\)\}\{S\}\\left\(1\-\\frac\{S\}\{N\}\\right\)\}\.\(12\)

This expression shows that the excess buffer sizeM−kM\-kdecreases at the rateO⁡\(1/S\)O\(1/\\sqrt\{S\}\): larger samples produce more accurate thresholds and therefore require smaller buffers\.

### Optimal Sample Size

Beyond the mandatory linear scan through the input, Prof\-K incurs two additional costs: \(i\) sampling and selecting the threshold fromSSvalues, and \(ii\) exact top\-kkon a buffer of sizeMM\. We model total overhead as:

C=τA​S\+τB​M​log⁡k,C=\\tau\_\{A\}S\+\\tau\_\{B\}M\\log k,\(13\)whereτA\\tau\_\{A\}andτB\\tau\_\{B\}are hardware\-dependent constants capturing the per\-element cost of sample processing and of retaining the top\-kkelements in the candidate buffer, respectively\.

###### Theorem 5\(Optimal sample size\)\.

WhenS≪NS\\ll N\(sof≈1f\\approx 1\), the minimizer ofCCsatisfies

S∗=\(z​τB​log⁡k2​τA\)2/3​\(k⁡\(N−k\)\)1/3\.S^\{\*\}=\\left\(\\frac\{z\\tau\_\{B\}\\log k\}\{2\\tau\_\{A\}\}\\right\)^\{2/3\}\\big\(k\(N\-k\)\\big\)^\{1/3\}\.\(14\)For the common regimek≪Nk\\ll N, this simplifies toS∗∝\(k​N\)1/3S^\{\*\}\\propto\(kN\)^\{1/3\}\.

This scaling is noteworthy: the optimal sample size grows as\(k​N\)1/3\(kN\)^\{1/3\}, much more slowly than eitherkkorNNindividually\. For example, withN=109N=10^\{9\}andk=100k=100, we haveS∗≈4,600S^\{\*\}\\approx 4\{,\}600—a tiny fraction of the input\.

In practice, we restrictSSto the interval\[Smin,Smax\]\[S\_\{\\min\},S\_\{\\max\}\], whereSmin=32768S\_\{\\min\}=32768avoids high\-variance threshold estimates andSmaxS\_\{\\max\}is chosen so that sampling overhead remains negligible relative to the mandatoryO⁡\(N\)O\(N\)streaming pass \(e\.g\.,Smax=217=131,072S\_\{\\max\}=2^\{17\}=131\{,\}072in typical GPU implementations\)\. The ratioτB/τA\\tau\_\{B\}/\\tau\_\{A\}may be estimated by lightweight microbenchmarking; a default choice ofτB/τA=1\\tau\_\{B\}/\\tau\_\{A\}=1performs robustly across a broad range of\(N,k\)\(N,k\)\. Whenα​S<−ln⁡εA\\alpha S<\-\\ln\\varepsilon\_\{A\}\(very smallkkcombined with modest sample size\), the Gaussian approximation becomes unreliable\. In this regime, we conservatively sett=1t=1\(using the sample maximum as threshold\), accepting occasional fallback to exact top\-kkcomputation\.

### Parameter Selection

[Algorithm1](https://arxiv.org/html/2608.12573#alg1)summarizes the adaptive parameter selection for Prof\-K\. For additional discussion of the algorithmic details, see[AppendixD](https://arxiv.org/html/2608.12573#A4)\.

Algorithm 1Adaptive parameter selection for Prof\-K0:

NN\(tensor size\),

kk\(top\-

kk\),

ε\\varepsilon\(target failure budget, default

10−310^\{\-3\}\),

Smin,SmaxS\_\{\\min\},S\_\{\\max\}\(sample size bounds\),

τB/τA\\tau\_\{B\}/\\tau\_\{A\}\(cost ratio, default

11\)

0:

S,t,MS,t,M\(sample size, threshold rank, buffer capacity\)

1:

z←Φ−1​\(1−ε/2\)\+Φ−1​\(1−ε/2\)z\\leftarrow\\Phi^\{\-1\}\(1\-\\varepsilon/2\)\+\\Phi^\{\-1\}\(1\-\\varepsilon/2\)\{

≈6\.58\\approx 6\.58for

ε=10−3\\varepsilon=10^\{\-3\}\}

2:

α←k/N\\alpha\\leftarrow k/N
3:

S∗←\(z⁡\(τB/τA\)​log2​k2\)2/3⋅\(k⁡\(N−k\)\)1/3S^\{\*\}\\leftarrow\\left\(\\frac\{z\(\\tau\_\{B\}/\\tau\_\{A\}\)\\log\_\{2\}k\}\{2\}\\right\)^\{2/3\}\\cdot\(k\(N\-k\)\)^\{1/3\}
4:

S←min⁡\(Smax,max⁡\(Smin,⌊S∗⌉\)\)S\\leftarrow\\min\(S\_\{\\max\},\\max\(S\_\{\\min\},\\lfloor S^\{\*\}\\rceil\)\)
5:

f←1−S/Nf\\leftarrow 1\-S/N
6:

t←max⁡\(1,⌈S​α\+z2​S​α​\(1−α\)​f⌉\)t\\leftarrow\\max\\\!\\big\(1,\\;\\lceil S\\alpha\+\\frac\{z\}\{2\}\\sqrt\{S\\alpha\(1\-\\alpha\)f\}\\,\\rceil\\big\)
7:

M←⌈k\+z​k⁡\(N−k\)S​f⌉M\\leftarrow\\big\\lceil k\+z\\sqrt\{\\frac\{k\(N\-k\)\}\{S\}f\}\\,\\big\\rceil
8:if

α​S<−ln⁡\(ε/2\)\\alpha S<\-\\ln\(\\varepsilon/2\)then

t←1t\\leftarrow 1% fallback\-prone regime

9:return

S,t,MS,t,M

## 5Experiments

We evaluate Prof\-K on synthetic benchmarks spanning a wide range of input sizes and sparsity levels, comparing against PyTorchtopkand RadiK\. We report wall\-clock speedups over key hyperparameters\. Finally, we demonstrate practical impact by integrating Prof\-K into BatchTopK Sparse Autoencoder \(SAE\) training, where top\-kkselection is invoked at every optimization step over large flattened activation tensors\. This setting tests whether kernel\-level improvements translate to end\-to\-end training acceleration, and whether the probabilistic relaxation affects learned representations or reconstruction quality\.

### Experimental Setup

We evaluate our Triton\-based implementation on high\-end data center \(DGX H100\) and consumer\-grade \(RTX 3060\) GPUs, benchmarking against nativetorch\.topkand RadiK\[[9](https://arxiv.org/html/2608.12573#bib.bib4)\]\. Our synthetic workloads encompass tensor sizesN∈\[222,230\]N\\in\[2^\{22\},2^\{30\}\]and selection sizesK∈\[25,219\]K\\in\[2^\{5\},2^\{19\}\]\. We generate tensors using Uniform, Standard Normal, and heavy\-tailed Pareto distributions\. Crucially, because our algorithm uniformly samples indices rather than values, it is inherently distribution\-agnostic; performance remains strictly invariant even under extreme outliers\. For all exact top\-kkevaluations, hyperparameters are conservatively configured to guarantee 100% accuracy\.

### Algorithm Performance

Figure[2](https://arxiv.org/html/2608.12573#S5.F2)illustrates global latency speedups across the\(N,k\)\(N,k\)parameter space\. Our approach exhibits pronounced advantages in the memory\-bandwidth\-bound large\-NNregime across both hardware tiers\. By isolating tensor size \(Figure[3](https://arxiv.org/html/2608.12573#S5.F3)\), we observe that our filtering mechanism drastically flattens the latency scaling curve compared to the strict bandwidth limits governing baseline methods\. Finally, strong scaling tests at constant sizeN=29N=29\(Figure[5](https://arxiv.org/html/2608.12573#S5.F5)\) confirm that our method maintains robust speedups even as the absolute retrieval volume grows proportionally with the dataset\. Latency remains highly competitive for largekk\.

![Refer to caption](https://arxiv.org/html/2608.12573v1/heatmap_dgxh100_torch.png)

![Refer to caption](https://arxiv.org/html/2608.12573v1/heatmap_dgxh100_radik.png)

![Refer to caption](https://arxiv.org/html/2608.12573v1/heatmap_rtx3060_torch.png)

![Refer to caption](https://arxiv.org/html/2608.12573v1/heatmap_rtx3060_radik.png)

Figure 2:Average latency speedup of our method across the\(N,k\)\(N,k\)space\.Top:Speedups on high\-end GPU \(DGX H100\) vstorch\.topk\(left\), and RadiK \(right\)\.Bottom:Speedups on consumer tier GPU \(RTX 3060\) vstorch\.topk\(left\), and RadiK \(right\)\. Our approach exhibits substantial advantages in the large\-NNregime\.Figure 3:End\-to\-end latency \(ms, log scale\) as a function ofNN, averaged acrosskk\. Shaded regions denote standard deviations across different random input distributions\. Prof\-K achieves the best performance starting fromN=224N=2^\{24\}\.Scalability with Selection Size \(kk\) and Memory Footprint\.A key advantage of our algorithm is its memory efficiency\. Figure[5](https://arxiv.org/html/2608.12573#S5.F5)fixes a massive tensor size \(e\.g\.,N=229N=2^\{29\}for H100\) and varieskk\. The paired bar chart demonstrates that our auxiliary memory consumption is negligible compared to SOTA radix\-based methods which requireO⁡\(N\)O\(N\)additional space on the GPU\. This renders our approach impervious to Out\-Of\-Memory errors when the sizeNNbecomes extremely large\.

Figure 4:Latency plotted againstkkfor a large size of tensorN=29N=29on DGX H100\. Both top\-kkmethods scale well withkk\.![Refer to caption](https://arxiv.org/html/2608.12573v1/images/Code_Generated_Image.png)Figure 5:Auxiliary GPU memory overhead as a function ofkkfor a fixed tensor size \(N=229N=2^\{29\}\) \(logarithmic scale\)\. Our method’s memory footprint is bounded by a function of the buffer capacity, avoiding OOM errors on massive arrays\.

### Application: BatchTopK Sparse Autoencoders

To evaluate the practical impact of Prof\-K in a real machine learning workload, we replace the exact batch\-global top\-kkoperator inside BatchTopK Sparse Autoencoders \(SAEs\) with our probabilistic filtering procedure\. BatchTopK SAEs are a natural testbed because top\-kkselection is executed at every optimization step over a large flattened activation tensor, making it a recurring systems bottleneck during training\. This setting therefore tests whether faster top\-kkselection yields measurable end\-to\-end acceleration in a realistic sparse training pipeline rather than only in standalone synthetic benchmarks\.

We consider the BatchTopK setup of Bussmann et al\.\[[3](https://arxiv.org/html/2608.12573#bib.bib1)\], but focus on the computational cost of the selection primitive rather than introducing a new SAE objective\. In each training step, the post\-ReLU latent activations form a matrix of shapeB×ddictB\\times d\_\{\\text\{dict\}\}, which is flattened into a vector of lengthN=B⋅ddictN=B\\cdot d\_\{\\text\{dict\}\}\. BatchTopK then selects the largestK=B⋅kK=B\\cdot kactivations globally across the batch\. This matches the operating regime targeted by Prof\-K: very largeNN, relatively small retained fractionK/NK/N, and repeated invocation inside the training loop\.

In our experiment, we use activations fromEleutherAI/pythia\-70m\-dedupedat the MLP submodule of layer 1 and train BatchTopK SAEs on OpenWebText, following the same training hyperparameter regime as Bussmann et al\.\[[3](https://arxiv.org/html/2608.12573#bib.bib1)\], namely batch sizeB=4096B=4096, learning rate3×10−43\\times 10^\{\-4\}, and target sparsity levelk=32k=32\. We run both methods for30,00030\{,\}000training steps\. For the dictionary size, we useddict=12288d\_\{\\text\{dict\}\}=12288, which is one of the standard settings reported in their GPT\-2 Small experiments\. Thus,

k=32,ddict=12288,B=4096,k=32,\\qquad d\_\{\\text\{dict\}\}=12288,\\qquad B=4096,which gives

N=B⋅ddict=50,331,648,K=B⋅k=131,072\.N=B\\cdot d\_\{\\text\{dict\}\}=50\{,\}331\{,\}648,\\qquad K=B\\cdot k=131\{,\}072\.The experiment is conducted on an NVIDIA DGX A100 system\.

We compare a baseline implementation based on Torch𝚝𝚘𝚙𝚔\\mathtt\{topk\}against a Prof\-K variant that uses probabilistic filtering followed by exact refinement on the retained candidate buffer\. Both runs use the same optimizer, model activations, and SAE hyperparameters, differing only in the implementation of the batch\-global top\-kkroutine\.

The systems results are shown in[Figure6](https://arxiv.org/html/2608.12573#S5.F6)\. Prof\-K reduces mean top\-kktime from approximately2\.562\.56ms to1\.131\.13ms, corresponding to a2\.27×2\.27\\timeskernel\-level speedup\. At the full training\-step level, the effect is smaller but still clearly measurable: mean step time decreases from approximately30\.7330\.73ms to29\.4229\.42ms, corresponding to a1\.044×1\.044\\timesend\-to\-end speedup\. This translates into a total training\-time reduction of approximately4\.25%4\.25\\%\.

The quality results are shown in[Figure7](https://arxiv.org/html/2608.12573#S5.F7)\. We track normalized mean squared error \(NMSE\), cross\-entropy degradation, fraction of variance explained, and average active latents per sample\. Across these metrics, the Prof\-K and Torch𝚝𝚘𝚙𝚔\\mathtt\{topk\}curves are nearly indistinguishable: reconstruction quality improves at the same rate, sparsity remains matched to the intended operating point, and downstream degradation remains comparable throughout training\. This indicates that replacing exact selection with Prof\-K preserves the optimization behavior of the surrounding SAE training pipeline\.

Taken together, these results show that Prof\-K is useful beyond top\-kkmicrobenchmarks\. In BatchTopK SAE training, it reduces the cost of the top\-kkprimitive while preserving training dynamics and final reconstruction quality\. Although the resulting end\-to\-end speedup in this setting is modest, even a few percent reduction in per\-step runtime can translate into hours of wall\-clock savings for larger language models or longer\-running training workloads\.

![Refer to caption](https://arxiv.org/html/2608.12573v1/images/speed_variant_10_paper_candidate.png)Figure 6:Runtime comparison in BatchTopK SAE training\.Left:mean top\-kkkernel runtime for Torch𝚝𝚘𝚙𝚔\\mathtt\{topk\}and Prof\-K\.Right:reduction in total training time induced by Prof\-K, reported relative to Torch𝚝𝚘𝚙𝚔\\mathtt\{topk\}\. Prof\-K substantially reduces the cost of the top\-kkprimitive and yields a smaller but consistent improvement in overall wall\-clock training time\.![Refer to caption](https://arxiv.org/html/2608.12573v1/images/ffff.png)Figure 7:BatchTopK SAE training quality comparison\. We compare Torchtopkand Prof\-K in terms of NMSEa and cross\-entropy degradation\[[3](https://arxiv.org/html/2608.12573#bib.bib1)\], as well as fraction of variance explained and average number of active latents per sample\[[10](https://arxiv.org/html/2608.12573#bib.bib15)\]\. The curves remain closely matched throughout training, indicating that replacing exact top\-kkwith Prof\-K preserves reconstruction quality, downstream behavior, and effective sparsity\.

## 6Conclusions

This work proposes Prof\-K, a probabilistic one\-pass filtering method for efficient top\-kkselection fromNNinput elements, achieving the largest gains in the large\-scale\-NN, small\-to\-moderate\-kkregime\. By estimating a conservative threshold from a small random sample, Prof\-K reduces the cost of exact selection while providing distribution\-agnostic correctness guarantees\. Our theoretical analysis establishes explicit bounds on recall failure and buffer overflow probabilities, together with an approximately optimal sample size scaling as\(k​N\)1/3\(kN\)^\{1/3\}\. Experiments on synthetic benchmarks and BatchTopK SAE training demonstrate substantial reductions in top\-kklatency while preserving downstream training quality and sparsity behavior\. These results show that probabilistic filtering can serve as an effective and robust systems\-level optimization for large\-scale sparse machine learning workloads\.

Limitations of the Prof\-K algorithm are discussed in Appendix[A](https://arxiv.org/html/2608.12573#A1)\.

## References

- \[1\]R\. Akbarinia, E\. Pacitti, and P\. Valduriez\(2011\)Best position algorithms for efficient top\-k query processing\.Information Systems36\(6\),pp\. 973–989\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§2](https://arxiv.org/html/2608.12573#S2.p1.1)\.
- \[2\]T\. Alabi, J\. D\. Blanchard, B\. Gordon, and R\. Steinbach\(2012\)Fast k\-selection algorithms for graphics processing units\.Journal of Experimental Algorithmics \(JEA\)17,pp\. 4–1\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p2.1),[§2](https://arxiv.org/html/2608.12573#S2.p2.1)\.
- \[3\]B\. Bussmann, P\. Leask, and N\. Nanda\(2024\)Batchtopk sparse autoencoders\.arXiv preprint arXiv:2412\.06410\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§1](https://arxiv.org/html/2608.12573#S1.p5.1),[§2](https://arxiv.org/html/2608.12573#S2.p4.1),[Figure 7](https://arxiv.org/html/2608.12573#S5.F7),[§5](https://arxiv.org/html/2608.12573#S5.SS0.SSS0.Px3.p2.1),[§5](https://arxiv.org/html/2608.12573#S5.SS0.SSS0.Px3.p3.1)\.
- \[4\]H\. A\. David and H\. N\. Nagaraja\(2004\)Order statistics\.John Wiley & Sons\.Cited by:[Appendix C](https://arxiv.org/html/2608.12573#A3.SS0.SSS0.Px2.p1.1)\.
- \[5\]W\. Fedus, B\. Zoph, and N\. Shazeer\(2022\)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research23\(120\),pp\. 1–39\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§2](https://arxiv.org/html/2608.12573#S2.p4.1)\.
- \[6\]A\. Gaihre, D\. Zheng, S\. Weitze, L\. Li, S\. L\. Song, C\. Ding, X\. S\. Li, and H\. Liu\(2021\)Dr\. top\-k: delegate\-centric top\-k on gpus\.InProceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis,pp\. 1–14\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p2.1),[§2](https://arxiv.org/html/2608.12573#S2.p2.1)\.
- \[7\]M\. Greenwald and S\. Khanna\(2001\)Space\-efficient online computation of quantile summaries\.ACM SIGMOD Record30\(2\),pp\. 58–66\.Cited by:[§2](https://arxiv.org/html/2608.12573#S2.p3.1)\.
- \[8\]D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. Chen\(2020\)Gshard: scaling giant models with conditional computation and automatic sharding\.arXiv preprint arXiv:2006\.16668\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§2](https://arxiv.org/html/2608.12573#S2.p4.1)\.
- \[9\]Y\. Li, B\. Zhou, J\. Zhang, X\. Wei, Y\. Li, and Y\. Chen\(2024\)RadiK: scalable and optimized gpu\-parallel radix top\-k selection\.InProceedings of the 38th ACM International Conference on Supercomputing,pp\. 537–548\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p2.1),[§1](https://arxiv.org/html/2608.12573#S1.p4.1),[§2](https://arxiv.org/html/2608.12573#S2.p2.1),[§5](https://arxiv.org/html/2608.12573#S5.SS0.SSS0.Px1.p1.1)\.
- \[10\]S\. Marks, A\. Karvonen, and A\. Mueller\(2024\)Dictionary\_learning\.Note:[https://github\.com/saprmarks/dictionary\_learning](https://github.com/saprmarks/dictionary_learning)Cited by:[Figure 7](https://arxiv.org/html/2608.12573#S5.F7)\.
- \[11\]A\. Radford, J\. Wu, R\. Child, D\. Luan, D\. Amodei, I\. Sutskever,et al\.\(2019\)Language models are unsupervised multitask learners\.OpenAI blog1\(8\),pp\. 9\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p5.1)\.
- \[12\]A\. Shanbhag, H\. Pirk, and S\. Madden\(2018\)Efficient top\-k query processing on massively parallel hardware\.InProceedings of the 2018 International Conference on Management of Data,pp\. 1557–1570\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§2](https://arxiv.org/html/2608.12573#S2.p1.1)\.
- \[13\]N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. Dean\(2017\)The sparsely\-gated mixture\-of\-experts layer\.Outrageously large neural networks2,pp\. 2\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p1.1),[§2](https://arxiv.org/html/2608.12573#S2.p4.1)\.
- \[14\]G\. Team, M\. Riviere, S\. Pathak, P\. G\. Sessa, C\. Hardin, S\. Bhupatiraju, L\. Hussenot, T\. Mesnard, B\. Shahriari, A\. Ramé,et al\.\(2024\)Gemma 2: improving open language models at a practical size\.arXiv preprint arXiv:2408\.00118\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p5.1)\.
- \[15\]X\. Xie, Y\. Luo, H\. Peng, and C\. Ding\(2024\)RTop\-k: ultra\-fast row\-wise top\-k selection for neural network acceleration on gpus\.InThe Thirteenth International Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p2.1),[§2](https://arxiv.org/html/2608.12573#S2.p2.1)\.
- \[16\]J\. Zhang, A\. Naruse, X\. Li, and Y\. Wang\(2023\)Parallel top\-k algorithms on gpu: a comprehensive study and new methods\.InProceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis,pp\. 1–13\.Cited by:[§1](https://arxiv.org/html/2608.12573#S1.p2.1),[§2](https://arxiv.org/html/2608.12573#S2.p2.1)\.

## Appendix ALimitations

The primary limitation of Prof\-K is that its benefits become less pronounced for smaller input sizesNN, where the overhead of sampling and filtering provides less advantage over highly optimized exact top\-kkimplementations\. Additionally, the method relies on fallback to exact top\-kkin rare cases where the probabilistic filter becomes unreliable or the candidate buffer overflows, introducing occasional overhead\. Our experiments are currently limited to single\-node GPU settings and sparse learning workloads based primarily on BatchTopK SAEs\. Moreover, additional care is required when applying a batched version of Prof\-K, as the probability of observing at least one failure increases with batch size\. Extending the approach to distributed environments and broader machine learning workloads remains an important direction for future work\.

## Appendix BImpact Statement and Declaration of LLM Usage

We do not anticipate direct negative societal impacts from this work\. The proposed method is a systems\-level optimization for efficient top\-kkselection and does not introduce new capabilities beyond standard machine learning functionality\. Any broader impact is indirect and arises from improved efficiency of existing models and workloads\.

LLMs were used exclusively for language editing and text refinement to improve the clarity of the manuscript\.

## Appendix CProofs

### Proof of[Lemma1](https://arxiv.org/html/2608.12573#Thmtheorem1)

The event\{R=r\}\\\{R=r\\\}requires: \(i\) exactlyt−1t\-1sampled indices among the topr−1r\-1positions, \(ii\) the element at rankrritself to be sampled, and \(iii\) the remainingS−tS\-tsampled indices to come from ranks belowrr\. Counting such configurations and dividing by\(NS\)\\binom\{N\}\{S\}yields[Equation1](https://arxiv.org/html/2608.12573#S4.E1)\.[Equation2](https://arxiv.org/html/2608.12573#S4.E2), which give the mean and variance, follows from standard properties of the negative hypergeometric distribution\.

### Proof of[Lemma2](https://arxiv.org/html/2608.12573#Thmtheorem2)

The result follows from the asymptotic normality of order statistics under simple random sampling without replacement\[[4](https://arxiv.org/html/2608.12573#bib.bib12)\]\. The variance includes the finite population correctionf=\(N−S\)/Nf=\(N\-S\)/N, which accounts for the reduced variability due to sampling without replacement whenSSis a non\-negligible fraction ofNN\.

### Proof of[Corollary4](https://arxiv.org/html/2608.12573#Thmtheorem4)

Substituting[Equations8](https://arxiv.org/html/2608.12573#S4.E8)and[11](https://arxiv.org/html/2608.12573#S4.E11)intoM=c​kM=ckgives

M=k⁡\[1\+zA\+zBS​1−αα​f\]=k\+\(zA\+zB\)​k2​\(1−α\)α​S​f\.M=k\\left\[1\+\\frac\{z\_\{A\}\+z\_\{B\}\}\{\\sqrt\{S\}\}\\sqrt\{\\frac\{1\-\\alpha\}\{\\alpha\}f\}\\right\]=k\+\(z\_\{A\}\+z\_\{B\}\)\\sqrt\{\\frac\{k^\{2\}\(1\-\\alpha\)\}\{\\alpha S\}f\}\.\(15\)Lettingz=zA\+zBz=z\_\{A\}\+z\_\{B\}and usingα=k/N\\alpha=k/N, we obtain

k2​\(1−α\)α=k2​\(N−k\)k=k⁡\(N−k\),\\frac\{k^\{2\}\(1\-\\alpha\)\}\{\\alpha\}=\\frac\{k^\{2\}\(N\-k\)\}\{k\}=k\(N\-k\),\(16\)which yields

M=k\+z​k⁡\(N−k\)S​f\.M=k\+z\\sqrt\{\\frac\{k\(N\-k\)\}\{S\}f\}\.\(17\)This proves the claim\.

### Proof of[Theorem5](https://arxiv.org/html/2608.12573#Thmtheorem5)

SubstitutingM=k\+z​k⁡\(N−k\)/SM=k\+z\\sqrt\{k\(N\-k\)/S\}into[Equation13](https://arxiv.org/html/2608.12573#S4.E13), differentiating with respect toSS, and setting the derivative to zero gives

d​Cd​S=τA−12τBlogk⋅zk⁡\(N−k\)S−3/2=0\.\\frac\{dC\}\{dS\}=\\tau\_\{A\}\-\\frac\{1\}\{2\}\\tau\_\{B\}\\log k\\cdot z\\sqrt\{k\(N\-k\)\}\\,S^\{\-3/2\}=0\.\(18\)Solving forSSyields the optimizer stated in[Equation14](https://arxiv.org/html/2608.12573#S4.E14)\. The second derivative is positive at this point, confirming that the solution is a minimum\. In the sparse regimek≪Nk\\ll N, we havek⁡\(N−k\)≈k​Nk\(N\-k\)\\approx kN, and thereforeS⋆∝\(k​N\)1/3S^\{\\star\}\\propto\(kN\)^\{1/3\}\.

## Appendix DAdditional Algorithmic Details

### Scaling withNNandkk\.

The candidate buffer grows sublinearly with both problem size and sparsity level:

M−k∝k⁡\(N−k\)S=O⁡\(k​NS\)\.M\-k\\propto\\sqrt\{\\frac\{k\(N\-k\)\}\{S\}\}=O\\\!\\left\(\\sqrt\{\\frac\{kN\}\{S\}\}\\right\)\.Using the optimal sampling rateS=Θ⁡\(\(k​N\)1/3\)S=\\Theta\(\(kN\)^\{1/3\}\)yields

M−k=O⁡\(\(k​N\)1/3\),M\-k=O\\\!\\left\(\(kN\)^\{1/3\}\\right\),which is asymptotically much smaller thanNNin the practically relevant regimek≪Nk\\ll N\. Thus, Prof\-K replaces an exact selection problem overNNelements with one over a substantially smaller candidate set\.

### Fallback as a Safety Mechanism

Rather than overprovisioningSSorMMto satisfy the target failure budget in pathological regimes \(which, for very smallα=k/N\\alpha=k/N, would requireS∝1/αS\\propto 1/\\alpha, Prof\-K reverts to exact top\-kkwhenever the probabilistic conditions are not met economically\. Because such events are rare by design, their amortized cost is negligible in practice\.

### Efficient GPU Implementation

All stages of Prof\-K, i\.e\., sampling, threshold estimation and filtering, the main compute payload, are amenable to parallelization\. In particular, the filtering scan can be efficiently implemented as a fused kernel in CUDA or in OpenAI’s Triton, enabling different parts of the large tensor to be processed simultaneously, saving only a fraction of the data to a pre\-allocated buffer\.

### Why Prof\-K Excels Whenkkis Small

Whenk≪Nk\\ll N, even a generous buffer multiplierc=M/kc=M/kcorresponds to a modest absolute buffer size\. For instance, withN=229N=2^\{29\}andk=64k=64, the buffer may contain only a few thousand elements \(M≈5,000M\\approx 5\{,\}000\), making the final exact top\-kkextremely cheap relative to scanning half a billion elements\.

### Distribution\-Agnostic Guarantees

Unlike quantile\-estimation methods that assume smooth or well\-behaved distributions, Prof\-K’s guarantees stem from combinatorial properties of random sampling\. The analysis depends only on ranks, not values, ensuring robustness against heavy\-tailed, multimodal, or adversarial inputs\.

### Graceful Degradation via Fallback

Rather than over\-provisioningSSorMMto handle pathological corner cases, Prof\-K simply falls back to exact top\-kkwhen needed\. If such events occur with probabilityε\\varepsilon, the amortized overhead remainsO⁡\(ε\)O\(\\varepsilon\), which is negligible for typical choices likeε=10−3\\varepsilon=10^\{\-3\}\.