论文

UBTree:通过 Unigram 和 Bigram 模型进行并行树起草以进行推测解码

UBTree: Parallel Tree Drafting via Unigram and Bigram Models for Speculative Decoding

模型推理投机采样

摘要

推测性解码通过在单个目标模型传递中验证多个草稿标记来加速语言模型推理。最近的并行牵伸器在前沿生产模型中取得了突破性的性能,但由于牵伸多样性不足,随着目标分布熵的增加,其有效性下降。为了在不牺牲并行性的情况下克服这一瓶颈,我们引入了 UBTree,这是一种并行绘图器,它将 Unigram 提议器与 Bigram 选择器结合起来构建绘图树。一元词提议器使用标准交叉熵目标进行训练,为每个位置独立生成候选标记,而轻量级二元词选择器则预测相邻候选对之间的转移分数。与提议者不同的是,选择器是使用高温数据上的重整化 KL 目标进行训练的。这种树原生训练将监督范围扩大到贪婪路径之外,鼓励合理的替代分支,从而提高在树验证期间接受额外词元的机会。在 Qwen3-4B 和 Qwen3-8B 的七个标准化基准测试中,UBTree 比自回归解码平均加速了 $5.84$-$6.94\times$,并且在所有 28 项比较中均优于 DARTree。生产规模评估进一步证明了 UBTree 相对于 DSpark 等前沿基线的优势。