Reward-Guided Speculative Decoding for Efficient LLM Reasoning

• Speculative Decoding的不足
• RSD概括
• 具体细节
• • RSD算法流程
  • token接受
  • RSD概率分布理论分析
• 其他

Speculative Decoding的不足

speculative decoding虽然可以加速LLM inference，其要求严格的无偏性，这在概率上和target model完全一致，理论上不会降低生成质量，但是由于严格无偏性的条件过于苛刻，这会导致draft model生成的一些高质量的token即使正确，由于分布差异或概率稍低而被拒绝，进而造成额外的target model的调用代价，这个问题在复杂推理中更为严重。

RSD概括

RSD的核心想法：放宽draft model 生成内容的接受条件，允许有偏接受策略。
具体而言，通过引入一个process reward model，以reward作为接受/拒绝的条件，而不是严格按照概率分布匹配。
这样做的好处也很显然，生成内容接受条件放宽必然带来更多token的接受，进而减少target model调用的次数，同时有process reward model的监督也能保证生成质量。

这张图可以分为上下两部分看，上半部分是传统Speculative Decoding，下半部分是论文提出的RSD
可以看到上半部分对于draft model生成的每一份draft，都需要调用target model进行验证，而下半部分RSD则是根据reward决定是否调用target model进行生成。

具体细节

RSD算法流程

token接受

token接受的reward阈值。
这里是通过一个$\omega (\cdot )$的权重函数，将reward映射到$[0,1]$，进而可以使用采样方法。
$$w(y_i\mid z_i)={\omega }_r(y_i\mid z_i)=\omega (r(y_i\mid z_i))$$

算法2，这里对于$\omega (\cdot )$允许多种实现方式，如阶跃函数或更平滑的函数
对于不同的权重函数作者也讨论了不同实现的好处，同时在table1中给出了不同的权重函数。

• 对于阶跃函数在理论上最优
• 对于其他函数则更加平滑

这里二元阶跃函数是最优的。

RSD概率分布理论分析

进行相关符号的规定：

较高的奖励值 $r(y_i\mid z_i)$ 表示，在给定输入 $x$ 和之前已生成的步骤 $y_{1:i-1}$ 的情况下，该模型输出与期望的响应更契合的可能性更大。
所以target model M的expect reward应该是大于draft model m
$${\mathbb{E}}_{y_i\sim {\mathbf{P}}_M}[r(y_i|z_i)]\ge {\mathbb{E}}_{y_i\sim {\mathbf{P}}_m}[r(y_i|z_i)],(1)$$

理论分析所提出方法的RSD的分布$P_{RSD}$，其由$P_m、P_M$结合
$${\mathbf{P}}_{\mathrm{RSD}}(y_i\mid z_i)=w(y_i\mid z_i){\mathbf{P}}_m(y_i\mid z_i)+v(y_i\mid z_i){\mathbf{P}}_M(y_i\mid z_i)$$
其中$w(y_i\mid z_i)$是一个权重函数根据draft model输出的reward动态调整，而$v(y_i\mid z_i)$是一个常量，这保证了target model永远都有一部分参与，不至于完全依赖draft model出现断层，换句话始终有target model来兜底。

其他

只有在target model显著大于PRM时这个方法的效果才会比较明显，因为这里虽然减少了target model的验证步，但是引入了新的开销，PRM获得reward。