Probabilistic Top-K Sampling

• Probabilistic Top-K sampling is a framework that selects the K most relevant elements by optimizing over the probability simplex with sparsity constraints.
• It employs convex optimization and divergence regularization to derive efficient support recovery methods applicable in language decoding and streaming analytics.
• Practical implementations use adaptive threshold sampling and efficient search strategies to ensure unbiased, scalable performance while addressing the combinatorial challenges of large K.

Probabilistic Top-K sampling encompasses a family of statistical and algorithmic frameworks that aim to select or sample the $K$ highest-valued, most probable, or most relevant elements from a collection or distribution, using randomness either in the construction of the top-$K$ set, the sampling within it, or both. This paradigm underpins diverse applications including LLM decoding, streaming heavy-hitter estimation, preference modeling, and subgraph mining, all unified by the need for scalable, tractable, and often unbiased estimators or samplers for large-scale or otherwise challenging domains.

1. Theoretical Foundations: Convex Optimisation and Divergences

The modern theoretical understanding of probabilistic Top-K sampling is grounded in the formalisation of the decoding step as an optimisation problem on the probability simplex, regularised to encode distributional constraints and practical preferences. The master problem is typically written as

$$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$

where $s \in \mathbb{R}^n$ is a score vector (logits, utilities, or counts), $p \in \Delta_n$ is a decoding distribution, and $\Omega(p)$ encodes constraints such as sparsity ($|\text{supp}(p)| \le K$), entropy, or quadratic penalties (Ji et al., 20 Feb 2026).

For the hard Top-K constraint,

$$\Omega_{\text{Top-K}}(p) = \begin{cases} 0, & |\text{supp}(p)| \leq K\ +\infty, & \text{otherwise} \end{cases}$$

yielding the classical solution uniform over the top-$K$ elements:

$$p^*_i = \begin{cases} \frac{1}{K}, & i \in S_K \ 0, & \text{otherwise} \end{cases}$$

where $K$0 indexes the top-$K$1 entries of $K$2.

Alternative formulations interpret the task as finding a sparse distribution $K$3 close to a model distribution $K$4 in a chosen Bregman divergence, regularised by an $K$5 penalty:

$K$6

Under KL divergence, this recovers the Top-K heuristic; with Tsallis or quadratic divergences, this generalises to new classes of decoders that interpolate between softmax, Top-K, and more complex selection rules (Noarov et al., 25 May 2025).

2. Algorithmic Frameworks: Efficient Sampling and Support Recovery

Despite the combinatorial nature of Top-K support and $K$7 regularisation, efficient algorithms are available for the most important cases:

• Support Selection: The optimal nonzero set for the sparse decoder is always the top-$K$8 entries by $K$9-value or $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$0-score (“greedy support”) (Noarov et al., 25 May 2025).
• Renormalisation: Within the Top-K support, $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$1 is renormalised, often uniformly or proportionally to original probabilities (for Bregman/softmax cases).
• Efficient Search: The cost function in $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$2 is discretely convex; binary search (combined with $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$3 sorting) finds the optimal $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$4 efficiently (Noarov et al., 25 May 2025).
• Streaming Heavy Hitters: Adaptive threshold sampling uses random priorities per item occurrence and dynamically shrinks the threshold $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$5, ensuring that heavy hitters are included with probability $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$6, and maintains unbiased Horvitz–Thompson estimators (Ting, 2017).

3. Unbiased Estimation and Statistical Guarantees

A principal advantage of probabilistic Top-K sampling in streaming and large-scale settings is its unbiasedness and controllable error rates:

• Adaptive threshold sampling yields estimators $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$7 that are unbiased for true frequencies $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$8, regardless of the arrival order (Ting, 2017).
• Miss Probability: For heavy hitters, the probability of exclusion decays exponentially in their count: $$\max_{p \in \Delta_n} \Big[\langle p, s \rangle - \lambda \Omega(p)\Big]$$9.
• Convergence of Priority-Queue Top-K: In subgraph mining, the rank order of expected supports from randomised samplers converges in probability to the true top-$s \in \mathbb{R}^n$0 as the number of samples grows, assuming adequate ergodicity and mixing (Saha et al., 2014).

4. Generalisations and Applications

Probabilistic Top-K sampling is instantiated across diverse problem domains with tailored generative or sampling mechanisms:

Application Domain	Core Method or Model	Key Guarantee/Property
LLM Decoding	Top-K simplex optimisation / Bregman divergence	Greedy support, efficient solution
Streaming Frequency Estimation	Adaptive threshold sampling, Horvitz–Thompson	Unbiased estimation, exponential tails
Preference/Ranking Modeling	Mallows model for Top-K prefixes, profile-based PRIM	Exact sampler for partial rankings
Frequent Subgraph Mining	MCMC on subgraph space + finite Top-K queue	Ergodicity, priority-queue convergence

• Generalized Mallows Top-$s \in \mathbb{R}^n$1 Model: Provides a probability distribution on Top-$s \in \mathbb{R}^n$2 lists using a symmetric dispersion parameter $s \in \mathbb{R}^n$3 and efficiently samples top-$s \in \mathbb{R}^n$4 prefixes via dynamic programming and prefix-insertion (PRIM), with guaranteed exactness and polynomial precomputation in $s \in \mathbb{R}^n$5 (Haddadan et al., 24 Oct 2025).
• Probabilistic Frequent Subgraph Mining (FS³): Uses Markov Chain Monte Carlo with edge-support–biased transitions, prioritizing the discovery of subgraphs with high expected support and storing the best candidates in a finite queue, yielding superior scalability and precision over deterministic enumeration under large $s \in \mathbb{R}^n$6 (Saha et al., 2014).

5. Comparative Analysis: Connections to Other Schemes

The unified optimisation lens clarifies the relationships between classic and modern sampling schemes:

• Softmax/Temperature Sampling: Softmax with temperature $s \in \mathbb{R}^n$7 corresponds to negative-entropy regularisation and has support over the full simplex; Top-K emerges in the hard-sparsity (cardinality) limit (Ji et al., 20 Feb 2026).
• Top-P/Nucleus Sampling: A quantile-based support relaxation, replaces Top-K’s cardinality constraint with a probability mass threshold; the renormalisation likewise takes place over the dynamic support set.
• Sparsemax and General Bregman Decoders: $s \in \mathbb{R}^n$8 or Tsallis-$s \in \mathbb{R}^n$9 divergence gives rise to decoders that interpolate between hard support and soft assignments, providing a spectrum between exact Top-K, softmax, and other sparsification strategies (Noarov et al., 25 May 2025).
• Comparison with Multinomial Logit/Plackett-Luce: In ranking, the Plackett–Luce model offers a multiplicative, IIA-admitting closed-form, but lacks the nuanced correlation structure and learning guarantees of Mallows-based Top-K schemes (Haddadan et al., 24 Oct 2025).

6. Practicality, Efficiency, and Limitations

The computational demands of probabilistic Top-K sampling frameworks vary with application and data modality:

• Time complexity for Top-K simplex optimisation and Bregman decoding is $p \in \Delta_n$0 per step, dominated by sorting and k-dimensional projections.
• Streaming sketching achieves $p \in \Delta_n$1 amortised update with $p \in \Delta_n$2 storage, adapting to unknown or changing data distributions (Ting, 2017).
• Markov chain–based subgraph mining scales linearly in subgraph size for sampling, with queue insertion/eviction kept manageable via aggressive filtering and lexicographic ordering; empirical studies show linear scaling in $p \in \Delta_n$3, outperforming deterministic algorithms at scale (Saha et al., 2014).
• Generalized Mallows Top-K sampling incurs $p \in \Delta_n$4 preprocessing, but allows $p \in \Delta_n$5 per-sample cost, tractable for moderate $p \in \Delta_n$6 (Haddadan et al., 24 Oct 2025).

The main limitations are the combinatorial explosion of support sets for very large $p \in \Delta_n$7, the reliance on ergodicity and fast mixing in MCMC-based approaches, and $p \in \Delta_n$8 scaling in the Mallows context for large $p \in \Delta_n$9, which may be prohibitive in some applications.

7. Synthesis and Outlook

Probabilistic Top-K sampling frameworks bring rigor, unification, and extensibility to the diverse heuristics historically used for selection, ranking, and sampling from complex distributions. By situating these problems as regularised optimisations on the probability simplex, recent advances enable principled derivation, efficient implementation, and the construction of new decoders with transparent structural and statistical properties (Ji et al., 20 Feb 2026, Noarov et al., 25 May 2025). Applications across streaming analytics, natural language generation, preference modelling, and graph mining benefit from unbiasedness, adaptivity, convergence guarantees, and tractable computation. Contemporary research continues to extend these paradigms through new divergences, adaptive schemes, efficient learning algorithms, and scalable implementations, cementing probabilistic Top-K sampling as a foundational tool for large-scale statistical inference and decision making.