TVP:用期望接受长度优化投机解码草稿模型训练
TVP: TV is Almost All You Need for Speculative Decoding
研究投机解码的可以看看:把草稿模型训练目标换成 TVP,接受长度比 CE 高 9.8%,比 DFlash2 配方高最多 12.9%。
投机解码中草稿 token 的接受概率等于 1 减去草稿分布与目标分布之间的总变差(TV)距离,但草稿模型通常用交叉熵(CE)或 KL 训练。论文发现微调时 TV 目标在匹配良好调参下可持平或超过多种已发表的草稿训练目标,从零训练时固定 TV 与 KL 混合也有同样效果。针对 TV 在未覆盖 token 上梯度消失、接受长度目标中前缀连乘让靠后位置饥饿这两个问题,作者推导出基于期望接受长度的 TVP 损失族,其幂参数在覆盖导向和接受导向之间插值。在 3 种草稿架构和 3 个目标模型(含长草稿和块验证)上,TVP 微调时接受长度比 CE 高 9.8%,比公开的 DFlash2 训练配方高 6.2% 到 12.9%。
TVP: TV is Almost All You Need for Speculative Decoding
Speculative decoding accelerates LLM inference with a drafter whose proposed tokens are verified by the target model in parallel. Under lossless speculative sampling, a draft token is accepted with probability one minus the total variation (TV) distance between the target and draft distributions, yet drafters are often trained with cross-entropy (CE) or KL. We first show that under matched and well-tuned training, TV matches or surpasses several published draft-training objectives when fine-tuning, and a fixed TV and KL mixture does the same when training from scratch. We identify two remaining problems, TV's vanishing gradients on undercovered tokens and the prefix products that starve later positions in acceptance-length objectives. To address both, we derive TVP, a loss family based on expected acceptance length, whose power parameter interpolates between coverage and acceptance oriented objectives. Annealing the power from zero to one gives the highest or tied-highest acceptance length across three drafter architectures and three target models, including longer drafts and block verification. TVP improves acceptance length over CE by 9.8% when fine-tuning, and over the public DFlash2 training recipes by 6.2 to 12.9%.