投机解码中高质量draft tokens不该被拒绝

💡ICLR25的8/8/10/5的高分工作，确实值得读一下！

论文 ：Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment
地址 ：https://openreview.net/forum?id=mtSSFiqW6y

作者：Wzl

动机

speculative sampling的机制（draft model生成多个draft tokens让target model验证），虽然保证了与target model的output一致（无损），但是，这也造成许多高质量的draft tokens因为与target model不对齐，即不符合target model的概率分布被拒绝。

我之前在做投机采样的工作中也和朋友讨论过这个问题，结果是投机采样的这种机制确实会限制draft model的发挥，即使draft model生成的token高质量（高质量的含义是对输出内容有保证的延续，从而获得更正确的回答），也会被拒绝，所以在设置draft一轮的token数的时候反而比较保守，一般在5左右，生成再长后面的也不容易被接受。

这篇工作从LLM-as-a-judge的角度受到启发，去处理target model验证过程中碰到的non-aligned的token，切入点非常准。

方法

增加一个“judge” head去判断token embedding是正确还是错误。

训练数据：作者制作了一个叫TokenCourt的数据集，标注了wrong answer中的错误token。

实验

主要针对draft-model为8B，target时70B/405B这样的setting，对比了标准SD，Medusa，Eagle2，效果还是有明显提升的。