TVP损失函数：投机解码中TV距离的潜力与改进

投机解码（Speculative Decoding）通过草稿模型快速生成候选token，再由目标模型并行验证，从而加速大语言模型推理。在无损投机采样下，草稿token被接受的概率等于1减去目标分布与草稿分布之间的全变差（TV）距离。然而，现有草稿模型多采用交叉熵（CE）或KL散度训练。一项新的研究重新审视了TV距离在草稿训练中的作用，并提出了改进的TVP损失函数，在多个架构上实现了接受长度的显著提升。

TV距离在草稿训练中的潜力

研究首先表明，在匹配且调优良好的训练条件下，TV距离在微调时能够匹配或超越多个已发表的草稿训练目标；当从零开始训练时，固定的TV与KL混合目标也能达到类似效果。这意味着TV距离本身是一个被低估的损失函数选择。

然而，TV距离存在两个问题：一是对覆盖不足的token梯度消失，二是接受长度目标中的前缀乘积会抑制后续位置的贡献。这两个问题限制了TV距离的进一步应用。

全变差距离（Total Variation Distance）是两个概率分布之间差异的一种度量，定义为两个分布在各点上的概率差绝对值之和的一半，取值在0到1之间。在投机采样中，它直接刻画了草稿分布与目标分布的不一致程度：若草稿分布与目标分布完全一致，TV距离为0，所有草稿token都被接受；TV距离越大，被拒绝的概率越高。因此，接受概率恰好等于1减去TV距离，这一性质使TV距离成为与投机解码接受率天然对应的优化目标，与交叉熵、KL散度等仅衡量分布差异而不直接对应接受率的损失形成对比。

TVP损失函数的设计

为了解决上述问题，研究者推导出TVP，一个基于期望接受长度的损失函数族。其幂参数能够在覆盖导向和接受导向的目标之间插值。通过将幂参数从零退火到一，TVP在三种草稿模型架构和三种目标模型上取得了最高或并列最高的接受长度，包括更长的草稿和块验证场景。

实验结果显示，微调时TVP相比CE将接受长度提升了9.8%，相比公开的DFlash2训练方案提升了6.2%到12.9%。这些改进表明，TVP能够更有效地训练草稿模型，从而提升投机解码的整体加速效果。

接受长度（acceptance length）指投机解码中每次草稿生成后被目标模型连续接受并采纳的token数量。由于目标模型每次前向验证可以并行检查多个草稿token，接受长度越长，单次验证所覆盖的有效token越多，端到端加速比越高。接受长度与接受率并非简单线性关系：它依赖于前缀token被逐一接受的概率乘积，因此后续位置只有在前面所有位置均被接受时才有贡献。这一前缀依赖结构正是TVP引入幂参数来平衡覆盖与接受导向的动因。

对投机解码实践的影响

投机解码的加速效果直接取决于草稿token的接受率。TVP通过优化接受长度目标，在多个模型组合上实现了更优的接受率，这为实际部署中提升推理速度提供了新的训练策略。研究者指出，TVP的幂参数退火策略简单有效，易于集成到现有训练流程中。

未来，随着更多草稿架构的出现，TVP有望成为草稿模型训练的标准损失函数之一。

背景补充

块验证（block verification）是投机解码的一种变体，指草稿模型一次性生成一个固定长度的token块，目标模型对该块内所有token并行计算概率并进行验证与拒绝采样。与逐token验证相比，块验证能更充分地利用目标模型的并行计算能力，但对草稿token之间的连贯性要求更高，因为块内任一位置被拒绝都会截断后续接受。TVP在块验证场景下同样取得最优接受长度，说明其损失设计能够适应不同的验证粒度。