Source-linked AI summary
Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers
Xiaotong Ji, Rasul Tutunov, Matthieu Zimmer, Haitham Bou-Ammar
TL;DR
The paper addresses the treatment of decoding as heuristic knob tuning by formulating it as regularised optimisation over the probability simplex. It derives classical decoders as special cases and introduces Best-of-K as a KL-anchored coverage objective for multi-sample generation. BoK improves accuracy across evaluated settings, including an 18.6% gain for Qwen2.5-Math-7B on MATH500 at high temperature.
Problem
Decoding is commonly framed as a disconnected collection of heuristics, leaving its relationship to principled optimisation conceptually unclear.
Method
The paper formulates decoding as simplex optimisation that trades model score against regularisers and constraints, then uses the framework to design KL-anchored Best-of-K coverage objectives.
Results
18.6%: BoK raises Qwen2.5-Math-7B accuracy on MATH500 at τ=0.9 from 53.0% (Base) to 71.6%.
Takeaways & Limitations
The optimisation framework unifies classical decoding strategies and provides a principled route for designing samplers targeted at multi-sample coverage.
Takeaways & Limitations
The paper focuses on per-step decoding and identifies sequence-level objectives and broader compute-aware utilities as future directions.
Abstract
from arXiv · showhide
Decoding sits between a language model 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.
1 Introduction
The paper reframes decoding as principled optimisation over token distributions rather than a collection of disconnected heuristics. It develops this framework and demonstrates it with Best-of-K sampling for multi-sample pipelines.
- Unified view: Decoding strategies can be understood as solutions to explicit optimisation problems with different regularisers and constraints.Greedy, Softmax, and sparsity-inducing decoders correspond to different implicit objectives.
- Unified view: The distribution-first formulation trades off high-scoring tokens against structural preferences such as smoothness, sparsity, and proximity to a reference distribution.A decoder first selects a token distribution, then samples from it or takes its mode.
- Unified view: The framework derives optimality conditions on the probability simplex, distinguishing tokens assigned positive probability from those pushed to zero.These KKT-style conditions provide a general way to recover decoder structure.
- Best-of-K: Best-of-K decoding introduces a KL-anchored coverage objective for self-consistency and reranking under a fixed K-sample budget.The objective targets capturing high-quality candidates rather than relying on empirical sampling heuristics.
- Best-of-K: 18.6%: BoK raises Qwen2.5-Math-7B accuracy on MATH500 at τ=0.9 from 53.0% (Base) to 71.6%.This exceeds Top-K at 56.2% by +15.4%.
- Best-of-K: 1s: Using 5 mirror-ascent steps per token adds about 1s on MATH500, while 2 steps increase accuracy from 64.4% to 69.6%.The reported runtime comparison is 16.88s versus 15.84s for 5 steps.
2 Decoding, Sampling and Optimisation
The paper formalises decoding as choosing a probability distribution over next tokens and solving a constrained optimisation problem on the simplex. Regularisers, constraints, and simplex optimality conditions determine the resulting decoder.
- Conceptual shift: Decoding algorithms differ by the objectives they optimise and the constraints they impose, rather than constituting unrelated procedures.The framework presents greedy, Softmax, Top-K, Top-P, and sparse decoders as special cases.
- Distributional formulation: At each step, the decoder selects an auxiliary distribution over vocabulary tokens before sampling a token or selecting its mode.This formulation covers both stochastic and deterministic decoding.
- Distributional formulation: The model supplies a real-valued score s_t(v) for every vocabulary token, which the optimisation uses to evaluate candidate distributions.These scores are logits or log-probabilities produced by the language model.
- Special cases: Greedy decoding is represented by a degenerate distribution concentrated on the highest-scoring token, whereas sampling uses a non-degenerate positive-entropy distribution.Temperature-based sampling favours larger scores while retaining stochasticity.
- Master objective: The master objective balances expected model score against a regulariser and support constraints, with λ controlling regularisation strength.The regulariser can encode diversity, sparsity, or stability, while C_t constrains the support.
- Optimality conditions: When the regulariser is convex, the resulting constrained problem admits clean optimality conditions, but the simplex constraints must be handled explicitly.The constraints shape which tokens are active or inactive at the solution.
3 LLM Decoding Strategies are Different Regularisers
The paper unifies common decoding strategies as optimisation problems over the probability simplex, where regularisers and support constraints determine the resulting decoder. This framework recovers greedy, softmax, Top-K, Top-P, and sparse decoding through shared optimality conditions.
- Common decoding strategies arise as special cases of one master optimisation problem with different regularisers, strengths, and support constraints.The framework treats decoding rules as choices of optimisation structure rather than unrelated procedures.
- With no regularisation, the optimum places all probability mass on the highest-scoring token or tokens, recovering greedy decoding.If the maximiser is unique, the solution is a degenerate distribution; ties allow any distribution supported on the maximisers.
- Negative Shannon entropy produces the temperature-controlled softmax distribution, with the regularisation weight acting as the sampling temperature.The entropy regulariser keeps the optimum in the simplex interior and yields a closed-form softmax solution.
- Top-K combines an entropy-regularised objective with a fixed support constraint that forces tokens outside the highest-scoring K candidates to zero probability.The resulting sampler is a softmax restricted to the selected Top-K set.
- Top-P replaces Top-K’s fixed cardinality with a cumulative-mass constraint, selecting the smallest context-dependent nucleus whose model probability mass reaches threshold p.Sampling then uses a temperature-controlled softmax renormalised over that nucleus; the nucleus expands under uncertainty and contracts under confidence.
- Replacing entropy with a quadratic penalty allows the solution to reach the simplex boundary and assign exactly zero probability to low-scoring candidates.This produces an adaptive truncation associated with Sparsemax-style decoding, avoiding entropy’s strictly positive tail probabilities.
4 Going Beyond Current Decoders
The paper extends its simplex-optimisation framework beyond closed-form decoders by introducing mirror ascent for objectives without analytic solutions, then applies it to Best-of-K sampling for multi-sample coverage.
- Mirror Ascent: Mirror ascent solves simplex optimisation when the objective is well-defined but lacks a closed-form solution.It preserves non-negativity and normalisation by construction.
- Mirror Ascent: Projected gradient ascent implicitly uses squared L2 geometry, which can mismatch probability-distribution behaviour near the simplex boundary.The paper associates this mismatch with unstable updates, aggressive redistribution, and repeated constraint interactions.
- Mirror Ascent: Replacing the L2 proximity term with a Bregman divergence generalises the optimisation to geometry induced by a convex potential, including simplex-compatible manifolds.Choosing the Euclidean norm recovers the projected-gradient formulation as a special case.
- Best-of-K Samplers: Best-of-K adds a coverage objective to the master problem for settings where K samples should contain high-quality alternatives rather than duplicate likely continuations.Its objective trades off model score against multi-sample coverage and is solved using mirror ascent.
- Best-of-K Samplers: BoK has an anti-collapse bias: it rewards allocating probability to valuable, under-covered tokens instead of repeatedly increasing already-likely tokens.This behaviour follows from the coverage objective's decreasing dependence on token probability for K > 1.
- Evaluation: +18.6% accuracy: BoK raises Qwen2.5-Math-7B on MATH500 at τ=0.9 from 53.0% to 71.6%, exceeding Top-K at 56.2% by +15.4%.The paper reports broader gains at τ=0.9 on GPQA and HumanEval, while near-deterministic settings can remain comparable or slightly favour Base or Top-K.
- Evaluation: 5 mirror-ascent steps per token add modest runtime relative to base decoding across the three benchmarks.Reported MATH500 runtime is 16.88s for BoK versus 15.84s for base decoding; HumanEval is faster under BoK in the reported implementation.
- Evaluation: 2 mirror-ascent steps improve MATH500 accuracy from 64.4% to 69.6% with negligible runtime increase, while additional steps beyond 5 yield marginal accuracy changes.Five steps reach 73.0% at 16.88s under the reported configuration.
5 Conclusion and Future Work
The paper unifies widely used decoding rules as solutions to one probability-simplex optimisation template and demonstrates its practical use with Best-of-K. It closes by proposing sequence-level objectives, compute-aware utilities, and richer constraint sets as future directions.
- Conclusion: Greedy, Softmax, Top-K, Top-P, and sparse decoders emerge from one master optimisation problem through different regularisers and feasible sets.Optimality conditions determine which tokens remain active or become inactive.
- Conclusion: Best-of-K introduces a regulariser that targets multi-sample coverage and can use the same mirror-ascent machinery.This demonstrates that the framework supports concrete decoder design beyond explaining existing methods.
- Future Work: Sequence-level objectives could couple decisions across time to enforce coverage, length, or style constraints globally rather than locally.The paper identifies this as a natural extension beyond its per-step focus.
- Future Work: Future work could explore compute-aware utilities for reranking, verifier selection, and self-consistency, alongside richer constraint sets beyond the simplex.Suggested constraints include structured sparsity, group constraints, and dynamic support sets linked to external tools or retrieval modules.