← 返回 2026-02-23

将解码视为概率单纯形上的优化:从Top-K到Top-P(核采样)再到Best-of-K采样器 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 📅 2026-02-20 👍 13 2026-07-13 08:35
优化理论 大语言模型 概率单纯形 解码策略 采样方法

提出统一优化框架解释所有解码策略,并设计新的Best-of-K采样器提升多样本生成质量

前置知识

概率单纯形 (Probability Simplex)

概率单纯形是所有合法概率分布构成的几何空间。对于词汇表V,概率单纯形∆(V)定义为所有满足∑q(v)=1且q(v)≥0的向量q的集合。在解码场景中,我们需要在每一步选择一个位于该单纯形上的分布q来决定下一个token的概率分配。单纯形的几何特性决定了优化问题的边界行为:当解落在内部时所有token都有正概率,当落在边界或顶点时某些token的概率为零。

理解概率单纯形是理解本文优化框架的基础,因为所有解码策略都被统一为在该空间上的优化问题,不同解码器的区别在于它们在单纯形上的不同位置和形状。

KKT条件 (Karush-Kuhn-Tucker Conditions)

KKT条件是约束优化问题最优解的必要条件。在本文中,由于优化变量q必须满足单纯形约束(概率非负且和为1),最优解必须满足两类条件:对于正概率的token(活跃集),梯度条件成立;对于零概率的token(非活跃集),梯度方向必须满足不等式约束。这些条件就像一把万能钥匙,一旦推导出来,代入不同的正则化器就能立即得到对应的解码器结构。

KKT条件是本文理论推导的核心工具,使得作者能够从统一框架出发,通过选择不同的正则化器和约束集,系统地推导出所有经典解码策略作为特例。

正则化 (Regularization)

正则化是在优化目标中加入的惩罚项,用于控制解的性质。在解码语境中,正则化器Ω(q)决定了我们希望解码分布具有的特性:负熵鼓励多样性、二次范数诱导稀疏性、KL散度锚定参考分布。正则化强度λ控制了模型得分与结构偏好之间的权衡。不同的正则化器产生不同的解码行为,这就是为什么解码方法不是启发式的,而是有原则的设计选择。

正则化是本文统一框架的核心概念,不同解码策略的本质区别不在于采样方式,而在于它们隐式优化的正则化目标不同。理解这一点就能理解论文的核心论点。

镜像上升法 (Mirror Ascent)

镜像上升法是梯度上升法在非欧几何空间上的推广。标准梯度上升使用L2距离衡量步长,但对概率分布而言这不是自然的距离度量。镜像上升用Bregman散度替代L2距离,当使用熵函数作为势函数时,Bregman散度就是KL散度。更新规则变为q_{j+1} ∝ q_j ⊙ exp(ηg_j),这种乘法形式天然保持概率分布的非负性和归一化,无需额外投影步骤。

当优化目标没有闭式解时(如Best-of-K),镜像上升提供了在单纯形上求解的原则性算法。它比投影梯度法更适合概率分布,因为KL几何比L2几何更匹配单纯形的结构。

多样本覆盖 (Multi-sample Coverage)

在自一致性、重排序等多样本生成场景中,我们不是生成单个答案而是生成K个样本。传统解码器为单次采样设计,容易重复采样高概率续写而忽略有价值的替代方案。覆盖度衡量的是:在K次独立采样中,某个token至少出现一次的概率,即1-(1-q(v))^K。这个量具有递减边际收益,意味着将概率分配给未被覆盖的token比增加已覆盖token的概率更有价值。

多样本覆盖是设计Best-of-K采样器的核心动机,它解决了当前多样本流水线的效率问题,通过显式优化覆盖度来提升K次采样的整体质量。

研究动机

当前大语言模型的解码策略被当作一系列独立的启发式方法:Top-K、温度采样、Top-P、贪心解码、束搜索等,每种都是基于经验或试错选择的'技巧'。这种观点使得解码与机器学习的其他部分在概念上脱节,无法系统地理解和设计新的解码器。更具体地说,在多样本生成场景(如自一致性、重排序、验证器选择)中,传统解码器为单次采样设计,会浪费预算反复采样相同的高概率续写,而对那些单独不太可能但集体有价值的替代方案分配不足。例如,在MATH500基准测试上,Qwen2.5-Math-7B在高温度τ=0.9时,标准采样准确率仅为53.0%,表明高多样性采样虽然探索性更强但可靠性大幅下降。

本文的目标是本文的目标是将解码策略重新理解为一个有原则的优化层:在每个token位置,我们在概率单纯形上求解一个正则化问题,平衡模型得分与结构偏好和约束。具体而言,作者希望:(1) 建立统一的理论框架,证明所有经典解码器都是单一优化模板的闭式解;(2) 提供一个生成性的设计工具,使研究者可以通过指定正则化器来设计具有期望行为的新解码器;(3) 利用该框架设计一个针对多样本场景优化的Best-of-K采样器,提升多样本生成的质量和效率。

与已有工作不同的是,本文的独特切入角度在于提出了'分布优先'的视角:解码器不必立即选择token,而是先选择一个分布,再从该分布采样或取模。这个看似微小的转变将解码转化为在概率单纯形上的干净优化问题。与以往工作不同,本文不仅推导了闭式解,还系统分析了最优性条件(KKT条件),这些条件就像一把万能钥匙,一旦推导出来,代入任何正则化器就能立即得到对应的解码器结构。此外,当闭式解不可得时,本文引入镜像上升法作为单纯形原生的求解器,为设计更丰富的解码目标提供了算法基础。这种从启发式到优化的视角转变,使得解码设计从经验艺术转变为可推导的科学。

核心方法

本文的方法可以分三个层次理解。首先,直觉上,作者观察到所有解码策略本质上都在做同一件事:在每个时间步选择一个概率分布来决定下一个token的概率分配。区别只在于这个分布是如何被'塑造'的——贪心解码把所有质量放在最高分token上,Softmax通过熵正则化产生平滑分布,Top-K/P通过约束支持集来裁剪词汇表。其次,技术路线上,作者将这个观察形式化为一个主优化问题:q* = arg max[⟨q, s_t⟩ - λΩ(q)] s.t. q ∈ C_t,其中第一项鼓励将概率放在高分token上,第二项通过正则化器控制分布形状,第三项通过约束集限制支持集。最后,通过选择不同的Ω和C_t,该框架恢复出所有经典解码器作为特例。

本文的核心创新点是将解码器的差异归结为正则化器的选择,而非采样机制的不同。具体而言:(1) 贪心解码对应λ=0(无正则化),解退化为顶点;(2) Softmax采样对应负熵正则化Ω(q)=∑q(v)log(q(v)),产生闭式softmax解,温度τ就是正则化强度λ;(3) Top-K/P通过约束支持集C_t实现,内部仍用熵正则化;(4) Sparsemax使用二次正则化Ω(q)=½||q||²,允许解落在单纯形边界产生稀疏分布。与已有方法的本质区别在于:以往每种解码器都被当作独立的启发式方法,而本文证明它们都是同一优化框架的不同实例,区别仅在于正则化器和约束的选择。这使得解码设计从'选择现成技巧'转变为'设计正则化器'。

方法步骤详情

本文方法包含三个主要步骤。第一步是理论推导:从主优化问题出发,使用KKT条件推导最优性条件。对于活跃token(q(v)>0),条件为s_t(v) - λ∂Ω/∂q(v) = η;对于非活跃token(q(v)=0),条件为s_t(v) - λ∂Ω/∂q(v) ≤ η。通过选择不同的Ω(q),代入这些条件即可得到对应的解码分布。例如,选择负熵正则化时,由于∂Ω/∂q(v)=1+log(q(v))在q(v)→0时趋向-∞,强制所有token保持活跃,最终推导出softmax分布。第二步是算法设计:当闭式解不可得时(如BoK),使用镜像上升法迭代求解。初始化q_0=p_t(模型分布),然后执行J步更新:q_{j+1} = q_j ⊙ exp(ηg_j) / ||q_j ⊙ exp(ηg_j)||_1,其中g_j是目标函数的梯度向量。第三步是Best-of-K设计:定义加权K-覆盖效用U_{K,t}(q) = ∑w_t(v)[1-(1-q(v))^K],结合KL锚定得到正则化器Ω^{BoK}(q) = KL(q||p_t) - βU_{K,t}(q),代入主问题后用镜像上升求解。

技术新颖性

本文的技术新颖性体现在多个层面。首先是理论统一性:虽然之前有工作注意到某些解码器与优化的联系,但本文首次将所有主要解码策略(贪心、Softmax、Top-K、Top-P、Sparsemax)统一为单一优化框架的特例,并严格推导了最优性条件。其次是几何视角:本文强调了单纯形几何的重要性,解释了为什么投影梯度法(隐式使用L2几何)不适合概率分布,而镜像上升(使用KL几何)才是正确选择。第三是生成性设计:框架不仅是描述性的,更是生成性的——作者展示了如何通过定义新的正则化器(BoK覆盖效用)来设计具有期望行为的新解码器。第四是递减收益性质:BoK使用的覆盖效用1-(1-q(v))^K具有递减边际收益,这提供了一种天然的'反坍缩'机制,鼓励将概率分配给未覆盖的token而非继续增加已覆盖token的概率。

解码即优化框架:主目标函数统一标准LLM解码策略
Figure 1: 解码即优化框架:主目标函数统一标准LLM解码策略

实验结果

本文在两个7B参数的Qwen模型(数学专用版Qwen2.5-Math-7B和通用版Qwen2.5-7B)和三个基准测试(MATH500、GPQA-diamond、HumanEval)上进行了全面评估。核心发现是:(1) BoK在高温度采样时带来最大提升——在Qwen2.5-Math-7B的MATH500上,τ=0.9时准确率从53.0%提升到71.6%(+18.6%),超过Top-K的56.2%达+15.4%;(2) 在GPQA-diamond上τ=0.9时提升+6.06%,在HumanEval上τ=0.9时提升+14.64%;(3) 改进在不同(β,λ)超参数设置下稳健,表明存在稳定的工作区域而非脆弱的调参结果;(4) 计算开销很小:使用5步镜像上升每token仅增加约1秒(MATH500上16.88s vs 15.84s);(5) 即使只用2步镜像上升,准确率也能从64.4%提升到69.6%,运行时增加可忽略(15.87s vs 15.84s)。对于通用模型Qwen2.5-7B,BoK同样有效:在MATH500上τ=0.9时从44.2%提升到60.2%,在HumanEval上τ=0.9时从45.12%提升到71.34%。

Qwen2.5-Math-7B在不同温度下的准确率
Table 1: Qwen2.5-Math-7B在不同温度下的准确率
Qwen2.5-7B(通用模型)在不同温度下的准确率
Table 2: Qwen2.5-7B(通用模型)在不同温度下的准确率
镜像梯度步数对MATH500准确率的影响
Table 3: 镜像梯度步数对MATH500准确率的影响
查看结构化数据
任务指标本文基线提升
MATH500 (数学推理) 准确率 (%) 71.6 53.0 (Base), 56.2 (Top-K) +18.6% vs Base, +15.4% vs Top-K (τ=0.9, Qwen2.5-Math-7B)
GPQA-diamond (科学问答) 准确率 (%) 36.36 30.30 (Base), 31.82 (Top-K) +6.06% vs Base, +4.54% vs Top-K (τ=0.9, Qwen2.5-Math-7B)
HumanEval (代码生成) 准确率 (%) 52.44 32.93 (Base), 37.80 (Top-K) +19.51% vs Base, +14.64% vs Top-K (τ=0.9, Qwen2.5-Math-7B)
MATH500 (通用模型) 准确率 (%) 60.2 44.2 (Base), 41.0 (Top-K) +16.0% vs Base, +19.2% vs Top-K (τ=0.9, Qwen2.5-7B)
HumanEval (通用模型) 准确率 (%) 69.51 45.12 (Base), 57.93 (Top-K) +24.39% vs Base, +11.58% vs Top-K (τ=0.9, Qwen2.5-7B)

局限与改进

本文的局限性包括:(1) 理论框架虽然统一了多种解码器,但对束搜索等序列级方法的覆盖有限,作者在结论中提到需要扩展到序列级目标;(2) Best-of-K的设计假设K次采样是独立同分布的,这在实际流水线中可能不完全成立(如带验证器的选择);(3) 实验主要在7B规模的模型上进行,对更大规模模型(如70B+)的适用性需要进一步验证;(4) 覆盖效用中的权重w_t(v)的选择目前较为粗糙(如使用模型得分的单调函数),更精细的设计可能带来进一步提升;(5) 镜像上升的收敛性分析较为初步,仅通过实验展示了快速收敛(2-5步即可),但缺乏理论收敛保证;(6) 论文未讨论BoK在流式生成或低延迟场景中的适用性,镜像上升的多次迭代可能在这些场景中成为瓶颈。

独立分析的弱点

本文的弱点可以从几个方面分析:(1) 覆盖效用的设计较为简单,当前使用1-(1-q(v))^K作为hit概率,未考虑token之间的语义关系——两个语义相近的token可能被重复计为'覆盖',而语义互补的token对可能更有价值;改进方向是引入语义感知的覆盖度量。(2) 权重w_t(v)的选择缺乏原则性指导,论文中使用模型得分的单调函数,但这可能不是最优选择;可以探索基于下游任务反馈的学习型权重。(3) 镜像上升的步长η和迭代次数J需要手动设定,虽然实验显示2-5步即可收敛,但缺乏自适应停止准则;改进方向是设计基于梯度范数或目标值变化的自适应策略。(4) 框架假设每个时间步的优化是独立的,未考虑跨时间步的耦合;对于需要全局一致性的任务(如风格控制),序列级优化可能更合适。(5) 实验评估主要关注准确率,未充分探索生成多样性、流畅性等其他维度。

未来方向

作者提出了几个有前景的未来方向:(1) 将框架扩展到序列级目标,耦合跨时间步的决策,实现全局覆盖、长度或风格约束;(2) 探索更广泛的计算感知目标,包括建模下游重排序、验证器选择或自一致性的效用函数;(3) 扩展到更丰富的约束集,如结构化稀疏、群组约束或依赖外部工具的动态支持集。基于本文成果还可以延伸:(4) 设计学习型正则化器,通过强化学习或元学习自动发现最优的Ω(q)形式;(5) 将BoK与其他解码技术(如推测解码、对比搜索)结合,探索协同效应;(6) 在多模态生成(图像、音频)中应用类似的优化视角;(7) 开发自适应BoK,根据模型置信度动态调整K和(β,λ)参数;(8) 理论分析镜像上升在BoK目标上的收敛速率和最优性间隙。

复现评估

论文的可复现性较好。作者使用了开源的Qwen2.5模型(7B规模),基准测试(MATH500、GPQA-diamond、HumanEval)都是公开可获取的标准化数据集。评估使用Qwen默认提示词和标准评估脚本,最大生成长度T_max=3072。BoK的实现基于清晰的算法描述(Algorithm 1),镜像上升更新有明确的数学公式。超参数设置(β=0.01/0.02/0.05, λ=0.1/0.2/0.5)在论文中详细列出,运行时开销也有具体数据(如MATH500上16.88s vs 15.84s)。然而,论文未明确说明是否开源代码,这可能影响完全复现。算力需求方面,7B模型在单GPU上即可运行,镜像上升每token仅需2-5步额外计算,总体开销增加约6%(16.88s/15.84s),适合大多数研究实验室。