Speculative decoding 投机采样原理和验证

最新推荐文章于 2024-08-30 20:51:37 发布

adam-liu

最新推荐文章于 2024-08-30 20:51:37 发布

阅读量1.8k

点赞数 13

文章标签：神经网络数学建模人工智能语言模型算法

本文链接：https://blog.csdn.net/qq_41664845/article/details/86470705

版权

概念

做法有很多，这里介绍最广泛被人认识的一种（朴素投机采样）

瓶颈：

大模型推理自回归采样，逐步串行解码。

生成的每个Token都需要将所有参数从存储单元传输到计算单元。

因此：内存访问带宽成为重要的瓶颈

大模型的推理有2个阶段，prefill（prompt的推理，可并行）和decode（generater过程，不可并行）

前者速度可以1k-50k token/sec，后者速度往往只有10-500 token/sec，最大相差两个数量级。

制约decode速度的最大瓶颈就是GEMV计算访存比太低了

现有的推理速度优化有很多办法，比如：in-flight batching、prefix cache等。

结果：内存访问得到了缓解，但带宽瓶颈依然存在。

投机采样（speculative decoding）

---- 一种解决内存访问带宽的根本性的方法

定义如下：

给定2个生成式模型：Target和Draft，它们共享相同的词汇表，但具有不同的参数量P1和P2，保证P1远远大于P2。
在相同的数据上训练，保证Draft和Target，在生成任务上的结果近似。
由于参数量的关系，Draft模型一次迭代（即生成一个token）的时间，Target模型可以生成多个token。

推理流程：

给定输入为： $X = \{x_{1},x_{2},x_{3}...x_{m}\}$ ，由Draft模型生成 $k$ 个token得到 $Y = \{y_{1},y_{2},y_{3}...x_{k}\}$
拼接输入 $X$ 和Draft模型输出 $Y$ ，得到 $Z = \{x_{1},x_{2},x_{3}...x_{m},y_{1},y_{2},y_{3}..y_{k}\}$ ，由Target模型对 $Z$ 进行推理
评估Target模型在 $\{y_{1},y_{2},y_{3},...,y_{k}\}$ 每个位置上的概率分布
选取保留前 $n\in\{1,2,3,...,k\}$ 个大模型Target、小模型Draft输出一致的token，作为一次迭代的输出

如下图所示：

结果
- 最好的情况
  - 小模型生成的 $k$ 个token和大模型都是一致的。
  - 那么小模型推理 $k$ 次，大模型推理一次，就有 $k+1$ 个可用的token。
  - 结果：节省了 $k$ 次大模型的推理
- 相对差的情况
  - 小模型生成的 $k$ 个token和大模型都不是一致的。
  - 那么小模型推理 $k$ 次，大模型推理一次，只有1个可用的token。
  - 结果：和原始推理性能相当
也就是说：
- 然后，在小模型足够近似的情况下，可以大幅度的提升推理性能。
- 投机采样完全不影响大模型的实际输出结果（推理精度）。

相关实验

Draft模型的选取4B，Target大模型选取70B。 https://arxiv.org/pdf/2302.01318

加速明显，尤其在HumanEval的加速几乎达到了2.5倍。

随着K的增大，考察固定128个token的总耗时，会出现先降再增的现象，XSum数据集上最优值是k=3。分析如下

如中图，随着k的增大，接受概率会逐渐降低
如右图，随着k的增大，平均每次调用的时长会增加，原因是小模型每次需要decode的K变大了。但此时大模型步长变大，当接受概率没有明显下降的时候，最终时长会出现下降，如左图
编程数据集Human Eval上的接受率要明显高于XSum，推测是因为编程的字符搭配更加有规律，更好预测