Foundations of Top-$k$ Decoding For Language Models

• URL: http://arxiv.org/abs/2505.19371v1
• Date: Sun, 25 May 2025 23:46:34 GMT
• Title: Foundations of Top-$k$ Decoding For Language Models
• Authors: Georgy Noarov, Soham Mallick, Tao Wang, Sunay Joshi, Yan Sun, Yangxinyu Xie, Mengxin Yu, Edgar Dobriban,
• Abstract summary: We develop a theoretical framework that both explains and generalizes top-$k$ decoding.<n>We show how to optimize it efficiently for a large class of divergences.
• Score: 19.73575905188064
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Top-$k$ decoding is a widely used method for sampling from LLMs: at each token, only the largest $k$ next-token-probabilities are kept, and the next token is sampled after re-normalizing them to sum to unity. Top-$k$ and other sampling methods are motivated by the intuition that true next-token distributions are sparse, and the noisy LLM probabilities need to be truncated. However, to our knowledge, a precise theoretical motivation for the use of top-$k$ decoding is missing. In this work, we develop a theoretical framework that both explains and generalizes top-$k$ decoding. We view decoding at a fixed token as the recovery of a sparse probability distribution. We consider \emph{Bregman decoders} obtained by minimizing a separable Bregman divergence (for both the \emph{primal} and \emph{dual} cases) with a sparsity-inducing $\ell_0$ regularization. Despite the combinatorial nature of the objective, we show how to optimize it efficiently for a large class of divergences. We show that the optimal decoding strategies are greedy, and further that the loss function is discretely convex in $k$, so that binary search provably and efficiently finds the optimal $k$. We show that top-$k$ decoding arises as a special case for the KL divergence, and identify new decoding strategies that have distinct behaviors (e.g., non-linearly up-weighting larger probabilities after re-normalization).

Related papers

• On Next-Token Prediction in LLMs: How End Goals Determine the Consistency of Decoding Algorithms [17.98959620987217]
  Next-token prediction trained using cross-entropy loss is the basis of most large language models.<n>This paper examines a few of these algorithms and studies their consistency with respect to various goals encoded as loss functions.
  arXiv  Detail & Related papers  (2025-05-16T12:38:45Z)
• Simple and Provable Scaling Laws for the Test-Time Compute of Large Language Models [70.07661254213181]
  We propose two algorithms that enjoy provable scaling laws for the test-time compute of large language models.<n>One is a two-stage knockout-style algorithm, where each candidate is evaluated by its average win rate against multiple opponents.<n>The other is a two-stage league-style algorithm, where each candidate is evaluated by its average win rate against multiple opponents.
  arXiv  Detail & Related papers  (2024-11-29T05:29:47Z)
• Theoretical limits of descending $\ell_0$ sparse-regression ML algorithms [0.0]
  We develop a generic analytical program for studying performance of the emphmaximum-likelihood (ML) decoding. Key ML performance parameter, the residual emphroot mean square error ($textbfRMSE$) is uncovered to exhibit the so-called emphphase-transition (PT) phenomenon. Concrete implementation and practical relevance of the Fl RDT typically rely on the ability to conduct a sizeable set of the underlying numerical evaluations.
  arXiv  Detail & Related papers  (2024-10-10T06:33:41Z)
• Projection by Convolution: Optimal Sample Complexity for Reinforcement Learning in Continuous-Space MDPs [56.237917407785545]
  We consider the problem of learning an $varepsilon$-optimal policy in a general class of continuous-space Markov decision processes (MDPs) having smooth Bellman operators. Key to our solution is a novel projection technique based on ideas from harmonic analysis. Our result bridges the gap between two popular but conflicting perspectives on continuous-space MDPs.
  arXiv  Detail & Related papers  (2024-05-10T09:58:47Z)
• Fast UCB-type algorithms for stochastic bandits with heavy and super heavy symmetric noise [45.60098988395789]
  We propose a new algorithm for constructing UCB-type algorithms for multi-armed bandits. We show that in the case of symmetric noise in the reward, we can achieve an $O(log TsqrtKTlog T)$ regret bound instead of $Oleft.
  arXiv  Detail & Related papers  (2024-02-10T22:38:21Z)
• SpecTr: Fast Speculative Decoding via Optimal Transport [30.18181671899423]
  We develop a new autoregressive sampling algorithm called $textitSpecTr$, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.
  arXiv  Detail & Related papers  (2023-10-23T17:47:34Z)
• Truncation Sampling as Language Model Desmoothing [115.28983143361681]
  Long samples of text from neural language models can be of poor quality. Truncation sampling algorithms set some words' probabilities to zero at each step. We introduce $eta$-sampling, which truncates words below an entropy-dependent probability threshold.
  arXiv  Detail & Related papers  (2022-10-27T05:52:35Z)
• Best Policy Identification in Linear MDPs [70.57916977441262]
  We investigate the problem of best identification in discounted linear Markov+Delta Decision in the fixed confidence setting under a generative model. The lower bound as the solution of an intricate non- optimization program can be used as the starting point to devise such algorithms.
  arXiv  Detail & Related papers  (2022-08-11T04:12:50Z)
• Approximate Function Evaluation via Multi-Armed Bandits [51.146684847667125]
  We study the problem of estimating the value of a known smooth function $f$ at an unknown point $boldsymbolmu in mathbbRn$, where each component $mu_i$ can be sampled via a noisy oracle. We design an instance-adaptive algorithm that learns to sample according to the importance of each coordinate, and with probability at least $1-delta$ returns an $epsilon$ accurate estimate of $f(boldsymbolmu)$.
  arXiv  Detail & Related papers  (2022-03-18T18:50:52Z)
• Towards Sample-Optimal Compressive Phase Retrieval with Sparse and Generative Priors [59.33977545294148]
  We show that $O(k log L)$ samples suffice to guarantee that the signal is close to any vector that minimizes an amplitude-based empirical loss function. We adapt this result to sparse phase retrieval, and show that $O(s log n)$ samples are sufficient for a similar guarantee when the underlying signal is $s$-sparse and $n$-dimensional.
  arXiv  Detail & Related papers  (2021-06-29T12:49:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.