DropMatch · Training-free Dropout Sampling for Semantic Token Acceptance in Speculative Decoding

1 · 动机:为什么 token 级精确匹配是浪费

1.1 历史脉络 — 从 lossless rejection 出发

Leviathan / Chen 2023 给的标准 SD 协议很优雅:逐位置 rejection sampling。给定 draft 分布 q 和 target 分布 p,以概率 min(1, p(x)/q(x)) 接受 draft 的 token x;否则按归一化后的残差 (p−q)₊ 重采样。这条规则数学上保证最终输出严格服从 p,所以叫 lossless。

但这个机制有一个隐含浪费:它判等的对象是 token id。在自然语言里,"积分" 和 "integral"、"计算" 和 "计算出"、"The" 和 "the"、甚至 "0.5" 和 "0.50" 这种从语义上几乎等价的 token,如果 draft 选 A、target 选 B,SD 直接拒绝并丢弃后续整段——尽管两者后续生成的 sentence 大概率是同一个意思。

更致命的是:在 SD 里 一次拒绝 = 后面 L−i 个 draft token 全部作废。因此一个"语义无伤大雅"的 mismatch 把 acceptance length τ 从 8 拉回 1,带来的损失是非线性的。

1.2 别的方案为什么不够

这两年涌现了 lossy SD 一支,允许"语义等价"的 token 通过。但每条路线都有自己的代价:

一个关键观察:上面所有"训 judge"的方案都在 OOD 时掉点。论文 Figure 1 直观地把这点画出来了——Auto-Judge 在数学训的 head 跑 IFEval 时,τ 高但 accuracy 崩;EAGLE-3 在英语训的 draft 跑 KoMT-Bench 时,task 性能保得住但 τ 从 ~5 跌到 1.46(Table 10)。训出来的 judge 本身就是新的 OOD 风险源。

另见 MARS (2601.15498) 笔记 — 同样是 relaxed-acceptance 路线,MARS 的"语义"用 logit margin 度量,DropMatch 用 dropout 经验分布度量。两者哲学不同:MARS 信任 target 自己的 confidence 几何,DropMatch 信任 target 的 ensemble 鲁棒性。

1.3 为什么这事不平凡

"接受语义等价的 token"听起来天经地义,但拆开看一堆坑:

• 什么叫"语义"? 没有第二个模型,你只能用 target 自己的输出来定义它。直接用 embedding 距离要引入 sentence encoder;直接用 KL 距离对 cutoff 极敏感。DropMatch 的取巧:语义 ≜ "在我自己的不确定性范围内可能产生的 token"。
• 怎么生成多样性? 温度采样需要重复 forward(贵),beam 不适合 verify 阶段。MC dropout 提供 free 的 ensemble effect,但整网 dropout 会破坏 KV cache 对齐——下一步 prefix 的 hidden 不再唯一,SD 框架直接坏掉。本文的 trick:只在 LM head 上 dropout,前面 transformer blocks 一字不差,KV cache 完全保留。
• 开销控制。 LM head 只占 forward 0.05%(Llama-3.1-70B),所以 K=5 路 head 计算 + JS 比较加起来 1.64%。再大 K 就开始划不来,因为 JS 是 O(K) 次 |V|=128k 的 softmax 比较。
• distortion 不能太大。 如果 p_drop=0.5,K 路头各说各话,接受规则会"无脑放行"。Table 1 显示 p=0.5 时 head Pass@1 已经从 81 掉到 65。所以 p_drop 的甜蜜区间约 0.2–0.3。
• 组合性。 必须能和 EAGLE-3 这种 tree-decoding + 自训 draft 的方案叠加。本文为此放弃 JS 改用 Naive Token-Matching(EAGLE-3 的 candidate 都来自同一个 dist,JS 计算冗余)。

2 · 背景速查

3 · 方法 §3.1 — LM head 上的 K 路 MC dropout

核心抽象:让 target 用 同一个 hidden state h_t 经过 K 个独立 dropout 掩码后通过 共享权重 W ∈ ℝ^|V|×d 的 LM head,产出 K 个 logits。

由 inverted-dropout, E[h⁽ⁱ⁾] = h, 所以 K 路 head 的"中心"和原始 head 在期望上一致。这是为什么不会偏离 target 太远——它是一个 controlled 的扰动而非 reparam。

为什么只挂在 LM head?

反向论证:如果整网 dropout, 第二个 token 的 hidden 也跟着扰动, KV cache 就要保 K 份, 显存翻 K 倍, 而且 SD 的 verify 阶段需要"batched 1 token" forward, dropout 在此时其实没意义。LM head 是"概率收口", dropout 在这里直接产生分布扰动, 物理上对应"target 在 logit 空间附近的小邻域", 是定义"语义等价"的最便宜锚点。

经验合理性

论文 Table 1 在 HumanEval 上分别测 K=5 个 head 的 Pass@1:

读法:p ≤ 0.3 时单个 dropout-head 仍然胜过 8B draft (Pass@1 72.6),说明这些 head 的输出仍是 high-quality target 分布,只是位置略有抖动。0.5 已经掉到不可用——这是 p_drop 的天花板。

4 · 方法 §3.2 — 双层接受判据

4.1 Naive Token-Matching

最简单的判据:让 K 路 head 各取 argmax 得到 token 集合 {y⁽¹⁾, ..., y^(K)}, draft token ŷ 只要在这个集合里就接受。

开销最低(只 K 次 argmax), 但有缺陷:它只看 top-1, 当 target 分布扁平、K 路 head 互相分散时, 可能会"放行"概率极低的 draft token。

4.2 JS-Divergence-Based

更稳的版本:先把 K 路 logits 平均后 softmax 得 centroid p̄(注意是 logit-mean 不是 prob-mean), 然后比较 draft 分布到 centroid 的 JS 与 K 路自身到 centroid 的最大 JS:

直觉:把"K 个 head 自己离 centroid 多远"当作"语义簇半径", 若 draft 离 centroid 不超过这个半径, 就承认它是这个簇的一员。

4.3 边界 case — 当 K 路高度集中

当 head 之间几乎一样(target 极自信), max_i JS ≈ 0, JS-only 判据会变得过于严苛(任何细微差异都被拒)。论文给的解决方法很务实:再加一条 majority 短路。

Table 2 显示 K=5 时 98.4% 的位置 5 个 head 全部投同一票——这种位置 JS 几乎一定要求 draft 严格相等, 但 majority rule 直接放行。Table 3 的数字也对得上:concentrated case 下 JS(c, draft)=0.0097 远大于 JS(c, heads)=0.0014, 单 JS 会拒, majority 救回来。

Input: ŷ_t (draft token), q̂_t (draft prob), {p^(i)_t}_{i=1..K}
1. p̄_t = softmax( mean_i l^(i)_t )         # logit 平均, 再 softmax
2. y^(i)_t = argmax p^(i)_t                # K 路 top-1
3. if JS(q̂_t ‖ p̄_t) ≤ max_i JS(p^(i)_t ‖ p̄_t):
       accept                              # 在簇半径内
4. elif majority({y^(i)_t}) == ŷ_t:
       accept                              # K 路投票一致放行
5. else: reject

5 · Worked example — 一个具体 token 走通整条流水

设当前 prefix 是 "求 ∫ x dx 的", draft (Llama-3.1-8B) 提议下一个 token 是 "积分"(token id 12345)。target (Llama-3.1-70B) 跑到这个位置后输出 hidden h_t ∈ ℝ⁸¹⁹²。我们做 K=5 路 dropout, p_drop=0.3:

Case A — Naive 判据

ŷ = "积分" ∈ {head1, head3, head5}, 立即 accept。开销:K 次 argmax, 微秒级。

Case B — JS 判据

centroid p̄ = softmax(mean of 5 logits) → 大约 0.55 集中在 "积分", 0.18 在 "integral", 其它分散。

• JS(q̂ ‖ p̄) ≈ 0.012 (draft 也是 "积分", 但分布没那么 sharp)
• max_i JS(p⁽ⁱ⁾ ‖ p̄): head4 离 centroid 最远, JS ≈ 0.041
• 0.012 ≤ 0.041 → accept

Case C — Vanilla SD 在同一位置怎么走?

Vanilla 算 p(积分)=0.55 / q(积分)=0.62 → accept prob = min(1, 0.55/0.62) = 0.887, 还有 11% 概率 reject; 即便 reject 后 resample 到 (p−q)₊, 后续 draft tokens 全部废掉。DropMatch 在这种 q 略乐观但语义对的位置直接 100% accept。

Case D — 一个 DropMatch 也该拒的位置

设 draft 提议 "结果"(token id 67890), 5 路 head argmax 全部是 "积分"。Naive: ŷ ∉ {y⁽ⁱ⁾} → reject。JS: q̂ 把质量集中在 "结果" 上, 与 centroid 极远, JS(q̂ ‖ p̄) ≫ max_i JS, reject。majority: 5 路 majority = "积分" ≠ ŷ, reject。三层都拦下来。

6 · 公式拆解 + 超参数敏感度

6.1 为什么 logit 平均而非 prob 平均?

centroid 选择有两种自然定义: (a) p̄ = (1/K) Σ p⁽ⁱ⁾ 直接平均概率; (b) p̄ = softmax((1/K) Σ l⁽ⁱ⁾) 平均 logit 再 softmax。论文采用 (b)。直觉:dropout 扰动是在 hidden 上做的线性变换 W·h, logit 比 prob 是更线性的"自然空间"; 平均 logit 等价于平均 dropout 后的 W·h, 与原始 W·E[h] 在期望意义上对齐。

6.2 K 与 p_drop 的 trade-off (理解性数字)

6.3 一个反向论证 — 为什么不能去掉 majority 这一条?

把 majority 条款拆掉, Table 3 显示 concentrated case 占 98.4% 的位置, 其中 JS(c, draft)=0.0097 比 JS(c, heads)=0.0014 大 7×; 单 JS 会把这些位置全部当作 reject。意味着没有 majority 你会在最容易 accept 的地方 (target 极 confident 处) 反而拒得最多, τ 立刻塌掉。

7 · 实验关键结果

7.1 Overhead — 这才是说服力最强的数字

实测 Llama-3.1-70B, batch=1, n_input=5。"Total fwd 反而更小"是噪声 — 重点是即使带上 K=5 + JS, 总开销 1.64% 几乎可以忽略, 这就为后续的 1.1× 加速打开了空间。

7.2 与 Standard SD (Table 6 节选)

读法:

• L=10 时 SD 自己经常变慢 (0.86x, 0.72x), 因为 draft 错一个就废后面 9 个, expected wasted work 高。DropMatch 在大 L 处收益最大, 因为它把"边缘 reject"救回来了。
• HumanEval 几乎没收益(Pass@1 反而掉一点)。原因:code 任务 head 全部一致(Table 2 的 98.4% majority), τ 已经接近上限, 没多少余地; 而代码语义对 token-level 等价非常敏感, distortion 容忍度更低。
• accuracy 普遍掉 0.4–1.0 个百分点, 在 lossy SD 容忍范围内。

7.3 与 EAGLE-3 叠加 (Table 7 节选)

关键洞察:EAGLE-3 自己在 L=7→9 处已经饱和(τ 不再涨, 4.65→4.68)。DropMatch 把这个饱和顶破——它击中的是 EAGLE-3 也无能为力的 OOD-like / 语义等价位置。叠加后 GSM8K accuracy 还涨 (78.75 → 80.00 at L=9), 这是 dropout 带来的 ensemble averaging 副作用。

7.4 OOD 健壮性 (Table 10 · KoMT)

EAGLE-3 在韩文上 τ 跌穿地板(draft model 没见过这种分布), 但 DropMatch 因为不依赖任何训出来的 component, τ 反而比 standard SD 还高一点。这是论文最有说服力的论点之一:training-free 在分布漂移下天然更稳。

8 · 与同类工作对比

核心差异化:

• vs MARS — 都是 lossy 且 training-free, 但 MARS 用 target 单次 forward 的 logit 几何; DropMatch 用 K 次 head dropout 的 ensemble 几何。后者对"target 在哪里有不确定性"刻画更直接, 代价是 K 倍 head 计算(但 head 太便宜)。
• vs Judge / Auto-Judge — 不需要训, 不会 OOD; 缺点是没有 judge 那种 task-specific 的可调 threshold(只能调 p_drop 间接控)。
• vs Vanilla — 用同样的 KV cache, 同样的 forward, 只是把 accept 函数换了; integration cost 几乎为零, 这是它能立刻叠在所有人头上的关键工程优势。

9 · 局限 / 个人 take / 待验证问题

• 没有严格 distortion bound. 论文承认是 lossy, 但没给"输出分布与 target 之间 TV / KL 距离的上界", 仅靠 p_drop<0.3 的经验性甜蜜点。这点比 vanilla SD 的 lossless 保证弱。
• 对 sharp 分布场景收益有限. HumanEval / 严格语法任务上几乎不加速, 因为 head 已 98.4% 一致, 没有"语义放行空间"。这是它与 code-domain workloads 不太合适的根本原因。
• K, p_drop 的最优值 task-dependent. 论文都用 K=5, p_drop=0.3, 但没系统扫 K∈{3,5,8,10} × p∈{0.1,0.2,0.3} 的二维 grid。在 reasoning vs IFEval vs code 上最优值很可能不同。
• JS 计算的实际 overhead 在大 batch 下可能放大. Table 4 是 batch=1 测的; batch=128 时 1.64% 的 head overhead 可能因为内存带宽 contention 变多。论文的 Table 5 显示 batch=128 还有 1.10x, 但这是 throughput 不是 head overhead 本身。
• 与 tree decoding 的耦合不优雅. EAGLE-3 时只能退回 Naive Token-Matching, JS 用不上。意味着语义边界判别能力在 tree-SD 上反而被砍。这部分论文一笔带过。
• 个人 take: 这个方法的贡献不在"新机制", 在于"重新框定问题":lossy SD 不一定要训新组件, 可以纯粹把 target 模型自身的 epistemic uncertainty 挖出来用。这是把 Bayesian DL 的老技术(MC dropout, 2016)套在 SD 接受判据上, 一种漂亮的 system-level 复用。

follow-up 待验证

1. p_drop 是否能 per-position 自适应? 比如根据 head entropy 决定开多大 dropout。
2. K 路是否可以异构(不同 p_drop, 或 head 上做 weight noise 而非 input noise)?
3. 替换 dropout 为 SWAG 或 deep ensemble (head-only) 会不会更稳但仍 cheap?
4. 能否给出最坏情况下输出分布与 target 的 TV bound? 至少能否在 small-vocab toy 模型上验证?
5. 大 batch 下 JS 计算的 attention-bound 行为如何? 在 vLLM 的 paged attention 上是否会破坏 batched verify 的 fused kernel?
6. 与 RL post-training rollout 的组合 — RL 训练里 sampling temperature 高, head ensemble 是否会和已有的 stochastic sampling 冗余, 还是互补?

10 · Memory points