推测解码:拒绝后怎么采样

黎 浩然/ 15 10 月, 2026/ 大语言模型/LARGELANGUAGEMODEL/LLM, 教程/TUTORIAL, 机器学习/MACHINELEARNING/ 0 comments

推测解码可以让小模型先写候选,再由目标模型验证。但“不改变目标分布”并不是靠一句“大模型把关”成立的:接受概率和拒绝后的补采样分布都必须正确。尤其是拒绝后直接按目标分布重抽,在常见的接受规则下会引入偏差。

文章目录
  1. 先限定“不损失”的含义
  2. 一个固定前缀下,接受了多少概率质量
  3. 三个token就能看出错误回退
  4. 先用精确枚举,不急着跑大量随机样本
  5. 边界情况决定实现是否可靠
  6. 多token草稿只能保留连续接受的前缀
  7. 温度与截断以后,才是实际要保持的分布
  8. 最高概率token相同,接受率仍可能不高
  9. 接受率高,也不自动等于端到端更快
  10. 来源与补发说明

下面用三个token的精确分数例子,把偏差算出来,再验证修正后的概率总量。Read in English。

先限定“不损失”的含义

这里讨论ICML 2023论文Fast Inference from Transformers via Speculative Decoding的精确采样思路。论文已经过会议发表,早于本次计划时段;这篇是算法与验收方法解读,不报道新发布。

所谓分布不变,指满足算法条件时,输出遵循目标采样分布。它不承诺某个随机种子下逐token文本与另一条执行路径完全相同,也不表示草稿模型的原始分布与目标模型相同。正确性约束的是最终采样机制。

一个固定前缀下,接受了多少概率质量

先只看固定前缀的下一个token。目标分布记p,草稿分布记q,二者针对同一个token空间且都已归一化。草稿先按q抽x,再以min(1,p(x)/q(x))的概率接受。只有q(x)>0的候选才会被抽到,因此该候选上的比值有定义。

token x通过接受路径到达最终输出的概率,等于q(x)乘接受概率,也就是min(p(x),q(x))。把它们求和得到整体接受概率A。目标还欠下的质量是p(x)−min(p(x),q(x))=max(p(x)−q(x),0)。

拒绝事件总概率是1−A,所以在拒绝之后,应把这部分尚未填补的质量除以1−A,作为条件补采样分布。接受路径填交集,拒绝路径补缺口,合起来才是p。

1
草稿按 q 提议
2
接受路径
保留 min(p,q)
3
拒绝路径
补上 max(p−q,0)
4
两条路径相加
最终质量等于 p
原创概率质量图:残差是尚未覆盖的目标概率,不是再复制一份完整目标分布。

三个token就能看出错误回退

构造目标p=(0.2,0.5,0.3),草稿q=(0.5,0.3,0.2)。接受路径留下(0.2,0.3,0.2),共0.7;剩余目标质量是(0,0.2,0.1),拒绝总概率为0.3。归一化后的补采样分布应是(0,2/3,1/3)。

如果拒绝后直接从完整p抽样,最终质量会是(0.2,0.3,0.2)+0.3×(0.2,0.5,0.3)=(0.26,0.45,0.29),不等于目标。错误在于已被接受路径覆盖的概率又被补了一遍,真正不足的部分却没有补齐。

1
目标 p
0.20 / 0.50 / 0.30
2
接受质量
0.20 / 0.30 / 0.20
3
正确补足
0.00 / 0.20 / 0.10
4
错误回退的最终结果
0.26 / 0.45 / 0.29
原创三token反例:所有数值为教学构造;最后一格是错误方法的总分布,不是实测生成频率。

先用精确枚举,不急着跑大量随机样本

这个离散例子可以直接枚举全部概率路径,不必用采样频率猜算法是否正确。Python的Fraction让0.2等小数以有理数精确表示,避免把舍入差当成分布偏差。

# Exact enumeration for one fixed prefix, not a decoder or benchmark.
from fractions import Fraction as F

def corrected(p,q):
    if len(p) != len(q) or not p:
        raise ValueError("matching nonempty vocabularies required")
    if any(x < 0 for x in p+q) or sum(p)!=1 or sum(q)!=1:
        raise ValueError("normalized nonnegative distributions required")
    kept = [min(a,b) for a,b in zip(p,q)]
    rejected = 1-sum(kept)
    if rejected == 0:
        return kept, None, kept
    residual = [max(F(0), a-b)/rejected for a,b in zip(p,q)]
    output = [a+rejected*r for a,r in zip(kept,residual)]
    return kept, residual, output

p = [F(2,10),F(5,10),F(3,10)]
q = [F(5,10),F(3,10),F(2,10)]
kept, residual, output = corrected(p,q)
assert kept == [F(2,10),F(3,10),F(2,10)]
assert residual == [F(0),F(2,3),F(1,3)]
assert output == p
reject = 1-sum(kept)
wrong = [a+reject*b for a,b in zip(kept,p)]
assert wrong == [F(26,100),F(45,100),F(29,100)]
assert wrong != p
assert corrected(p,p)[1] is None
assert corrected(p,[F(1),F(0),F(0)])[2] == p
assert corrected([F(0),F(1)],[F(1),F(0)])[2] == [F(0),F(1)]
print("exact output:", [str(x) for x in output])
print("incorrect fallback:", [str(x) for x in wrong])

代码已经执行。它验证正确输出逐项等于p、错误回退确实偏移,还覆盖p=q、草稿对某些token概率为零,以及目标与草稿支撑集完全不相交的情况。程序只枚举单个固定前缀,不运行语言模型,不测量速度。

边界情况决定实现是否可靠

当p=q时,A=1,没有拒绝事件,残差归一化分母为零;正确处理是根本不进入该分支,而不是人为加一个小数后继续抽样。例子返回None,明确表达没有需要使用的条件分布。

当q对某token给零概率、p却非零,草稿不会提出该token,但拒绝后的残差可以把它补回来。所以不能只在草稿提议过的token列表上构造最终概率向量。若实现截断了候选空间,需要确认截断后的实际算法仍满足自己的正确性定义。

浮点实现还需要检查归一化、非负值和极小残差。不能把数值上很小的拒绝概率随意当作零,同时继续宣称精确保持原分布。近似优化可以有价值,但要清楚标注它改变了什么。

多token草稿只能保留连续接受的前缀

草稿长度大于一时,每一步的p和q都对应此前候选组成的条件前缀。目标模型可以在一次验证计算中得到多处条件概率,但这些分布并不是互相独立、随意拼接的。

在第一次拒绝处补上新token之后,后面的旧草稿是基于已经被替换的前缀生成的,不能把它们当作同一条有效路径继续保留。实现应截断到连续接受的部分,再按更新后的真实前缀继续。若全部接受,原论文还会从相应目标条件分布再抽一个token。

缓存也应与保留下来的前缀一致。只修文本序列、却遗留被拒绝位置对应的状态,会让后续计算不再对应当前序列。本文的单步枚举验证不了这类状态管理;集成测试需要额外覆盖。

温度与截断以后,才是实际要保持的分布

如果生成策略用了温度、top-k或top-p,目标p应该是这些规则之后真正用于抽样的分布,而不是没有处理过的原始logits。草稿q也必须对应它实际提出候选的机制。把不同阶段的概率代入同一比值,推导就不再适用。

还要保证token事件的含义一致。相同文本在不同分词器中可能有不同token序列,不能仅按词表数组下标相等就比较概率。支持不同分词器的方案需要额外的对齐算法,本篇的同词表证明没有覆盖它。

固定随机种子后输出不同,未必说明分布错误;反过来,少量输出看起来相似,也不能证明分布正确。先核对精确小例子,再在足量统计实验中检查偏差,最后单独观察任务质量。

最高概率token相同,接受率仍可能不高

例如p=(0.5,0.3,0.2),q=(0.9,0.05,0.05),两者最高概率token相同,但本篇随机采样规则的接受概率只有0.5+0.05+0.05=0.6。只比较argmax是否相同,会丢掉分布中其余概率质量的信息。

这个数值是单个固定前缀的构造例子,不能直接外推成模型的整体接受率。评估草稿模型时,要在实际任务与生成策略下统计对应前缀上的接受行为。贪心解码具有自己的目标规则,不应和这里的随机采样接受率混为一谈。

接受率高,也不自动等于端到端更快

草稿提议与验证都有成本,候选越长还可能浪费更多被拒绝后的计算。接受率提供复用程度的信息,不能单独决定墙钟时间。设备、批大小、生成长度和缓存开销都应留在测量条件里。

可以先用单步反例验收采样修正,再检查多步前缀和缓存,最后比较同一目标模型、同一生成策略下的端到端时间。站内LLM背压文章讨论服务并发;单请求解码收益与整个服务的吞吐、尾延迟要分别观察。

正确概率、正确状态与实际性能,是三份不同的证据。若这个精确枚举例子帮你避开一次“看似无损”的采样错误,欢迎在文末自愿支持。

来源与补发说明

ICML 2023正式论文页;原文算法1与附录A.1。图、三token数字与分数枚举为原创教学内容,无模型或GPU基准。10月15日晚补发,文章日期保留当天19:00。

支持

如果文中的解释或示例帮你解决了问题、节省了时间,欢迎自愿打赏支持本站。

微信支持二维码,点击查看大图
微信
支付宝支持二维码,点击查看大图
支付宝
Buy Me a Coffee,点击支持本站
Buy Me a Coffee

二维码可点击放大。更多支持方式见支持页面。

Share this Post

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注

*
*