Speculative Decoding (推测解码) 简介

Introduction

• speculative decoding (推测解码)的思想先有Deepmind 2023.02 发布（Accelerating Large Language Model Decoding with Speculative Sampling），后Google research 另一拨人在2023.05 又发了一篇 （Fast Inference from Transformers via Speculative Decoding）。基本思想都一致：
  • tokens的生成有难易之分，对于一些”简单”的tokens (比如code生成中一些约定俗成的写法), 大小模型一致程度会比较高，因此可以用小模型快速生成，已提高inference的速度
    • i.e. 在code生成中，如果已经生成了”import numpy”, 很自然地下面的token大概率是” as np”。
  • 小模型 （draft model）快速推理decode K个tokens, 大模型（target model）利用prefill的形式一次feedforward判断是否接受小模型的生成的K个tokens。即draft model推测，target model审核。
  • 该方式用于decode token的分布是和target model生成token的分布是一致的。即大模型的能力能够完全保证。
  • 在Google相关的模型 （T5, Chinchilla） 实验上有2x的speed up。个人推测在code这种有很多约定俗成的编写规范的前提下，应该speedup会更好一些。

Methods

1. 用小模型p做自回归采样连续生成 K个tokens。
2. 把生成的K个tokens和前缀拼接一起送进大模型q执行一次forward。
3. 使用大、小模型normalzied logits (i.e. softmax归一化后得到的概率)结果做比对，如果发现某个token小模型生成的不好，重新采样这个token。重复步骤1。
4. 如果小模型生成结果都满意，则用大模型采样下一个token。重复步骤1。

第2步，将K个tokens和前缀拼成一起作为大模型输入，和自回归相比，尽管计算量一样，但是K个tokens可以同时参与计算，计算访存比显著提升。

第3步，如何评价一个token生成的不好？

• 如果q(x) > p(x)（q，p表示在大小模型采样概率，也就是logits归一化后的概率分布), 即小模型虽然采样概率小于大模型，但依然sample到了该token，则保留。否者认为虽然sample到了该token，但大模型在该token上的prob较低，需要重新sample。q(x)/p(x) （cond. <1）越小，接受的概率越小。
• 如果大模型拒绝了该token (拒绝概率 1- q(x)/p(x)),则以一个新的概率分布重新sample token
  • 新的概率分布为 norm(max(0, q(x)/p(x))) （norm公式详见paper）
• 以上方式能保证采样的分布和大模型直接生成的分布一致。证明见paper appendix: Modified Rejection Sampling recovers the target distribution.