Foundations of Top-$k$ Decoding For Language Models
Georgy Noarov, Soham Mallick, Tao Wang, Sunay Joshi, Yan Sun, Yangxinyu Xie, Mengxin Yu, Edgar Dobriban
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$---as well as other---sampling methods are motivated by the intuition that true next-token distributions are sparse, and the noisy LLM estimates 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 for a fixed token as the estimation of a sparse probability distribution. We consider estimators obtained by minimizing a separable Bregman divergence 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 (for both \emph{primal} and \emph{dual} minimization). 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 behavior (e.g., non-linearly up-weighting larger probabilities after re-normalization). Our work thus facilitates understanding and generalizing the widely popular top-$k$ decoding method for LLMs.