Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers

Abstract: Decoding sits between a LLM and everything we do with it, yet it is still treated as a heuristic knob-tuning exercise. We argue decoding should be understood as a principled optimisation layer: at each token, we solve a regularised problem over the probability simplex that trades off model score against structural preferences and constraints. This single template recovers greedy decoding, Softmax sampling, Top-K, Top-P, and Sparsemax-style sparsity as special cases, and explains their common structure through optimality conditions. More importantly, the framework makes it easy to invent new decoders without folklore. We demonstrate this by designing Best-of-K (BoK), a KL-anchored coverage objective aimed at multi-sample pipelines (self-consistency, reranking, verifier selection). BoK targets the probability of covering good alternatives within a fixed K-sample budget and improves empirical performance. We show that such samples can improve accuracy by, for example, +18.6% for Qwen2.5-Math-7B on MATH500 at high sampling temperatures.

• The paper introduces a unifying convex optimization framework that recasts diverse decoding algorithms as special cases through structured regularization and constraints.
• It employs KKT conditions and mirror ascent to derive classic decoders like greedy, softmax, Top-K, and Top-P, clarifying their underlying geometry.
• The Best-of-K sampler is proposed to enhance multi-sample coverage, achieving up to 18.6% accuracy gains on benchmarks with minimal computational overhead.

Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers

Introduction

This work reconceptualises LLM decoding as a principled convex optimisation process on the probability simplex, rather than as a disparate set of heuristic sampling algorithms. The central thesis is that commonly-used decoders—greedy, softmax/temperature sampling, Top-K, Top-P (nucleus), and sparse projection—are all special cases of a master optimisation problem. Within this framework, regularizers and constraint sets encode the structural preferences characteristic of each method, unifying their treatment through Karush-Kuhn-Tucker (KKT) optimality conditions. This formal perspective enables both theoretical understanding and systematic design of new decoders, as exemplified by the introduction of the Best-of-K (BoK) sampler.

Unification of Decoding via Regularised Optimisation

Decoding is framed as, at each step, selecting a distribution $q_t(\cdot)$ over tokens on the simplex, maximising a tradeoff:

$$q^*_t = \arg\max_{q \in \Delta(\mathcal{V})} \left[ \langle q, s_t \rangle - \lambda \Omega(q) \right],\quad \text{s.t. }q \in \mathcal{C}_t$$

where $s_t$ denotes the model token scores, $\Omega(q)$ is a regularization functional (e.g., negative entropy for Softmax), $\lambda$ modulates the regularisation strength, and $\mathcal{C}_t$ enforces structural support constraints (e.g., Top-K/P).

This "distribution-first" perspective abstracts both stochastic and deterministic decoders as varying $\Omega$ and $\mathcal{C}_t$, decoupling algorithmic surface differences from underlying objective geometry. The framework generalises to arbitrary convex regularisers and feasible sets, with decoding solutions defined constructively by KKT conditions.

Derivation of Classical Decoders as Objective Specialisations

The general formulation recovers canonical algorithms algebraically:

• Greedy Decoding: $\lambda = 0$, no regularisation, yields a degenerate distribution (all mass on the top token).
• Softmax Sampling: $\Omega(q) = -\mathcal{H}(q)$, provides the interior point softmax solution; temperature arises as $\lambda$.
• Top-K/Top-P (Nucleus) Sampling: Add respective support constraints on $q$; recovers softmax over the top $k$ tokens or the smallest set whose probability exceeds $p$.
• Sparsemax: Quadratic regulariser $\Omega(q) = \frac{1}{2} \lVert q \rVert_2^2$; solution is driven to simplex faces, generating sparse distributions.

For each, KKT conditions cleanly partition tokens into active/inactive sets according to the interaction between $s_t$, $\Omega'$, and constraint geometry.

Mirror Ascent as a Simplex-Native Solver

Projecting gradient updates onto the simplex implicitly imposes Euclidean ($L_2$) geometry, which is misaligned for probability distributions. Mirror ascent, in contrast, utilises Bregman divergences (notably KL for the negative entropy potential) and yields multiplicative natural-gradient-like updates. This property is essential for stability and convergence efficiency when working with distributions or enforcing non-negativity.

Given a general regularised objective, mirror ascent iterates:

$$q_{j+1}^{(i)} \propto q_j^{(i)} \exp \left( \eta \frac{\partial f(q_j)}{\partial q^{(i)}} \right)$$

and remains robust for regularisers with complex (non-closed form) gradients, as in coverage-based objectives.

Best-of-K (BoK) Sampler: Targeting Multi-Sample Coverage

The authors design BoK regularisation to improve self-consistency and reranking pipelines that require sampling $K$ candidates and selecting the best, e.g., via reranking or verification. Instead of focusing on single-sample likelihood, BoK explicitly optimizes for the event that good alternatives are covered at least once across $K$ samples. Concretely, for each token, its coverage probability (i.e., appearing at least once in $K$ i.i.d. draws) is $1 - (1-q(v))^K$.

BoK incorporates:

$$\Omega^{\mathrm{(BoK)}}_t(q) = \mathrm{KL}(q\|p_t) - \beta \sum_v w_t(v)(1-(1-q(v))^K)$$

where $w_t(v)$ indicates token importance (e.g., score-based), $\beta$ regulates the emphasis on coverage, and the KL term constrains $q_t^*$ close to $p_t = \text{Softmax}(s_t)$.

BoK decoding is then:

$$q_t^* = \arg\max_{q \in \Delta(\mathcal{V})} \left[\langle q, s_t\rangle - \lambda {\rm KL}(q||p_t) + \bar\beta \mathcal{U}_{K,t}(q)\right]$$

solved efficiently by mirror ascent, with closed-form gradients for the coverage term.

Empirical Results and Efficiency

Extensive evaluation on Qwen2.5-Math-7B and Qwen2.5-7B (MATH500, GPQA-diamond, HumanEval) demonstrates that:

• BoK samplers substantially improve accuracy at high temperatures (e.g., up to +18.6% absolute on MATH500 at $\tau=0.9$), outperforming both vanilla and Top-K sampling.
• Gains are most significant where exploration (diversity from stochasticity) would otherwise degrade reliability; BoK regularisation enables broader coverage of good hypotheses without collapse.
• BoK is robust to $(\beta, \lambda)$ hyperparameters and converges with as few as 2-5 mirror ascent steps per token (minimal runtime overhead, no external verifier).
• Improvements are observed not only in mean performance, but also in robustness across tasks and sampling regimes.

Implications and Future Directions

This optimisation-centric decoding framework transforms decoder design into explicit regulariser and constraint selection. Theoretical and practical implications include:

• Principled Decoding Design: Enables automatic derivation of new decoders aligned with application-specific structure (e.g., constraint satisfaction, diversity, coverage).
• Bridging Heuristics and Theory: Deconstructs folklore sampling algorithms into first-principles objective-driven processes, with explicit geometry and optimality.
• Extension to Structured and Global Objectives: The approach is readily extensible to sequence-level objectives and structured constraints, including those for style, coverage, external verification, and retrieval augmentation.
• Rich Multi-Sample Utilities: Facilitates formal integration of compute-aware objectives relevant for reranking, self-consistency, and verifier cascades.

The framework provides a foundation for methodical, principled decoding layer development, moving the field beyond heuristic aggregation to regulariser-driven algorithm generation.

Conclusion

By recasting LLM decoding as convex optimisation on the simplex, this work synthesises canonical decoding rules, clarifies their underlying regularisation structure, and equips researchers with systematic tools for decoder design. Mirror ascent supplies an efficient algorithmic substrate for closed-form and iterative decoders, and the BoK coverage regulariser offers empirical improvements for multi-sample regimes common in modern LLM applications. This paradigm enables both interpretability and practical improvements in LLM output, establishing a new blueprint for decoding algorithm development (2602.18292).