Fast Inference from Transformers via Speculative Decoding

“Fast Inference from Transformers via Speculative Decoding”由Yaniv Leviathan等人撰写。文章提出了推测解码算法，通过并行计算多个令牌加速自回归模型推理，在不改变模型架构、训练过程和输出分布的前提下，实现2 - 3倍的加速。

1. 背景与动机：大型自回归模型（如Transformers）推理速度慢，解码K个令牌需串行运行K次模型。现有加速推理的方法，如降低推理成本、自适应计算等，通常需要改变模型架构、训练过程或重新训练，且无法保证输出相同。本文基于部分推理步骤有难易之分，且大模型推理瓶颈常在于内存带宽和通信而非算术运算的观察，提出通过增加并行性加速推理，即推测解码算法。

2. 推测解码

   • 概述：使用更高效的近似模型$M_{q}$生成$\gamma$个完成结果，再用目标模型$M_{p}$并行评估这些猜测及其概率，接受能导致相同分布的猜测，并根据调整后的分布采样修正被拒绝的猜测或添加额外令牌。
   • 标准化采样：将多种采样方法统一为从调整后的概率分布进行标准采样，便于后续处理。
   • 推测采样：从$q(x)$采样，若$q(x)≤p(x)$则保留，否则以$1-\frac{p(x)}{q(x)}$的概率拒绝并从调整后的分布$p'(x)$重新采样，可证明采样结果服从$p(x)$分布。

3. 分析

   • 生成令牌数：假设$\beta$独立同分布，用$\alpha=E(\beta)$衡量$M_{q}$对$M_{p}$的近似程度，单次运行算法生成的令牌数是上限为$\gamma+1$、成功概率为$1-\alpha$的截尾几何变量，期望生成令牌数为$\frac{1-\alpha^{\gamma+1}}{1-\alpha}$。
   • 计算$\alpha$：定义$D_{LK}(p, q)$衡量分布差异，证明$\beta=1-D_{LK}(p, q)$，进而得到$\alpha=1-E(D_{LK}(p, q))=E(min (p, q))$。
   • 运行时间改进：在有足够计算资源支持并行的情况下，算法减少对目标模型的调用次数，考虑近似模型$M_{q}$的运行成本，得到算法在总运行时间上的改进因子为$\frac{1-\alpha^{\gamma+1}}{(1-\alpha)(\gamma c+1)}$。
   • 算术运算次数：算法并行运行$\gamma+1$次$M_{p}$，增加了并发算术运算次数，其总运算次数增加因子为$\frac{(1-\alpha)(\gamma \hat{c}+\gamma+1)}{1-\alpha^{\gamma+1}}$ ，但内存访问次数可能减少。
   • 选择$\gamma$：根据$c$和$\alpha$数值求解可得到最优$\gamma$ ，以最大化运行时间改进。预测$\beta$并动态调整$\gamma$可能进一步提升性能。
   • 近似模型：推测采样对近似模型选择无限制，实验中使用现成的较小Transformer模型，近似模型比目标模型小两个数量级时效果较好，此外还有零成本近似模型（如n - gram模型）等多种类型。

4. 实验

   • 运行时间改进：在T5 - XXL模型的英德翻译和文本摘要任务上进行实验，用T5 - small作为近似模型时加速比最高，分别在翻译任务和摘要任务上实现2.6X - 3.4X和2.3X - 3.1X的加速，与理论预测相符。
   • 经验$\alpha$值：测量不同任务、模型和采样方法下的$\alpha$值，发现近似模型比目标模型小几个数量级时，$\alpha$值在0.5 - 0.9之间，且调整后的分布越尖锐，$\alpha$值越高，简单的一元和二元模型也能产生不可忽视的$\alpha$值。

5. 相关工作：对比其他加速推理方法，如蒸馏、稀疏化、量化等通用方法，以及自适应计算方法，这些方法通常需改变模型架构和训练过程，且不保证输出相同。与类似的利用推测执行加速解码的方法相比，本文方法支持一般随机设置，对近似模型限制少。

6. 讨论：推测解码在有足够计算资源时可显著加速推理，但增加并发会导致算术运算次数增加。未来研究方向包括探索与束搜索的兼容性、设计定制近似模型、采用分层算法、动态调整近似模型和$\gamma$、应用不同分布变换，以及拓展到其他领域等。