投机采样（Speculative Decoding）原理与实验验证
投机采样（Speculative Decoding）原理与实验验证
投机采样（Speculative Decoding）是一种用于解决大模型推理中内存访问带宽瓶颈的技术。通过使用一个参数量较小的Draft模型和一个参数量较大的Target模型，可以在保持推理精度的同时显著提升推理速度。本文将详细介绍投机采样的原理、实现流程，并通过实验数据展示其效果。
概念
大模型推理中的自回归采样过程存在一个主要瓶颈：每个生成的Token都需要将所有参数从存储单元传输到计算单元，导致内存访问带宽成为重要的性能瓶颈。具体来说，大模型的推理分为两个阶段：prefill（prompt的推理，可并行）和decode（generater过程，不可并行）。前者的速度可以达到1k-50k token/sec，而后者往往只有10-500 token/sec，相差两个数量级。制约decode速度的最大瓶颈是GEMV计算的访存比过低。
尽管已有in-flight batching、prefix cache等优化方法，但内存访问带宽的瓶颈依然存在。投机采样（speculative decoding）提供了一种根本性的解决方案。
投机采样的定义
投机采样通过使用两个生成式模型来优化推理过程：一个参数量较大的Target模型和一个参数量较小的Draft模型。这两个模型共享相同的词汇表，但参数量不同（P1远大于P2）。通过在相同的数据上训练，确保Draft模型和Target模型在生成任务上的结果近似。
推理流程如下：
- 给定输入，由Draft模型生成K个token得到输出。
- 将输入和Draft模型的输出拼接，得到新的输入序列，由Target模型进行推理。
- 评估Target模型在每个位置上的概率分布。
- 选取保留前K个大模型Target和小模型Draft输出一致的token，作为一次迭代的输出。
最佳与最差情况分析
- 最佳情况：小模型生成的K个token和大模型完全一致。此时，小模型推理K次，大模型推理一次，就有K个可用的token，节省了K-1次大模型的推理。
- 最差情况：小模型生成的K个token和大模型都不一致。此时，小模型推理K次，大模型推理一次，只有1个可用的token，性能与原始推理相当。
在小模型足够近似的情况下，投机采样可以大幅度提升推理性能，且完全不影响大模型的实际输出结果（推理精度）。
实验验证
实验中选择4B参数量的Draft模型和70B参数量的Target大模型进行对比。实验结果表明，在HumanEval数据集上，投机采样的加速效果非常明显，几乎达到了2.5倍。
随着K值的增大，考察固定128个token的总耗时，会出现先降再增的现象。在XSum数据集上，最优值是k=3。分析如下：
- 如中图所示，随着k的增大，接受概率会逐渐降低。
- 如右图所示，随着k的增大，平均每次调用的时长会增加，原因是小模型每次需要decode的K变大了。但此时大模型步长变大，当接受概率没有明显下降的时候，最终时长会出现下降，如左图所示。
- 在编程数据集Human Eval上的接受率要明显高于XSum，推测是因为编程的字符搭配更加有规律，更好预测。