arXiv:2505.19371cs.AIcs.LG2025-05NeurIPS被引 13

为语言模型的top-k采样提供理论依据,发现最优k可通过二分查找确定。

Foundations of Top-$k$ Decoding For Language Models

论文配图:Foundations of Top-$k$ Decoding For Language Models
图 1 · 摘自论文原文
  • 将采样视为稀疏概率分布的恢复问题,用Bregman散度和l0正则化建模。
  • 证明最优k可高效通过二分搜索找到,且目标函数在k上离散凸。
  • 揭示top-k是KL散度下的特例,提出新策略可非线性放大高概率项。

Top-k解码是语言模型中常用的采样方法:每一步仅保留最大k个下一个词的概率,再归一化后采样。该方法基于真实词分布稀疏的直觉,但缺乏精确的理论支持。本文构建了统一的理论框架,将固定位置的解码视为稀疏概率分布的恢复。我们考虑通过最小化可分Bregman散度(包括原对偶情形)并引入ℓ₀正则化实现稀疏性。尽管目标函数具有组合性,我们仍证明其在一大类散度下可高效优化。结果表明最优解码策略为贪心,且损失函数关于k离散凸,因此可通过二分查找高效求得最优k。我们证明top-k解码是KL散度下的特例,并发现了具有不同行为的新解码策略(例如,在归一化后非线性地增强较大概率值)。

原文摘要 · Abstract (English)

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).

语言模型采样方法理论分析

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。