什么是推测解码(Speculative Decoding)?
推测解码是一种加速自回归模型推理的采样算法,由 arXiv 论文《Fast Inference from Transformers via Speculative Decoding》提出。它借助更高效的近似模型先并行猜出若干候选词元,再由大模型并行校验,从而在不改变输出分布的前提下减少串行解码步数。
推测解码(Speculative Decoding)是一种用于加速自回归模型(autoregressive model)推理的采样算法。按照提出该方法的论文《Fast Inference from Transformers via Speculative Decoding》的描述,它让大模型在不改变输出分布的前提下更快地完成采样,办法是借助一个更高效的近似模型先并行地猜出若干候选词元(token),再让目标大模型并行地对这些候选做校验,从而一次生成多个词元。它要解决的问题是:像 Transformer 这类大型自回归模型,解码 K 个词元需要串行运行模型 K 次,推理速度受限于这种逐词串行的过程。
为什么重要
自回归生成的基本约束是:第 t 个词元的概率分布依赖于前面所有已生成的词元,因此标准解码必须一个词元一个词元地串行推进。论文在摘要中把这一点直接点明——解码 K 个词元就要对模型做 K 次串行运行。对于参数量很大的模型,每一次前向计算本身就很昂贵,串行次数又无法通过批处理摊薄,于是单条序列的生成延迟居高不下。
在推测解码出现之前,常见的加速思路大体分两类。一类是改变模型本身,例如重新训练、蒸馏或修改架构,代价是需要额外的训练资源,并且往往不能直接套用到已有的现成模型上。另一类是改变采样结果,例如使用更激进的近似解码策略,用输出质量的下降换取速度。论文强调其方法的定位与这两类都不同:它不需要重新训练,也不需要改动模型架构,可以直接作用于现成的模型,并且不改变输出分布。这三点合在一起,构成了它被关注的主要理由。
工作机制
论文把方法的出发点归结为两个观察:其一,困难的语言建模任务中往往包含更容易的子任务,这些子任务可以被更高效的模型较好地近似;其二,借助推测执行(speculative execution)和一种新的采样方法,可以让大模型在这些近似模型的输出上并行运行,从而在保持分布不变的情况下实现精确解码。具体过程可以拆成以下几个环节。
- 近似模型先行猜测。用一个计算代价更低的近似模型,按照常规自回归方式连续生成若干个候选词元。这一步是串行的,但因为模型更小、更快,整体开销较低。
- 目标模型并行校验。把近似模型给出的这串候选词元一次性送入目标大模型,让大模型在一次前向计算中同时算出每个位置上自己的概率分布,而不是像标准解码那样为每个词元各跑一次。
- 按新采样方法接受或拒绝。论文提出了一种新的采样方法,依据近似模型与目标模型在各位置上的概率关系,逐个决定候选词元是被接受还是被拒绝。被接受的词元可以直接作为输出,因此一次校验有可能同时产出多个词元。
- 拒绝后回退并继续。一旦某个候选词元被拒绝,就从该位置起丢弃后续候选,由目标模型在该位置重新采样一个词元,然后进入下一轮猜测与校验。
这套流程的关键性质在于「精确」:论文称其采样方法使得最终输出的分布与直接用目标模型解码所得的分布一致,因此加速不以输出改变为代价。此外,由于近似模型只影响候选的提出效率、不影响最终分布,方法对近似模型的选择相对宽容——只要它足够快、且与目标模型的输出有足够重合度,就能带来收益。
典型例子
论文给出的实验对象是 T5-XXL。作者在 T5X 这一实现上做了对比,报告的结果是相对标准 T5X 实现获得 2 倍到 3 倍的加速,并且输出完全相同。论文同时说明,该方法可以加速已有的现成模型,无需重新训练或改动架构。该工作被 ICML 2023 接收为 Oral 报告。
需要说明的是,上述数字来自该论文自身的实验设定与实现对比,属于论文报告的结果,并不构成对其它模型、其它硬件或其它推理框架的普遍性能承诺。
边界与常见误解
第一,推测解码不改变输出分布,这一点容易被误读成「它一定会让生成变快」。实际上加速效果取决于近似模型提出的候选被接受的比例:如果近似模型与目标模型的判断经常不一致,候选频繁被拒,那么额外的猜测与校验开销就可能抵消收益,甚至带来负收益。论文的加速结论是在其特定模型组合与实现下得到的。
第二,它不是模型压缩或量化。推测解码不减小目标模型的规模,也不降低其数值精度,目标模型仍然要完整地参与每一次校验前向计算。它优化的是「需要串行运行多少次」,而不是「每次运行有多贵」。
第三,它不是对输出质量的近似折衷。与一些用质量换速度的近似解码策略不同,论文明确将「不改变输出分布」作为方法性质之一。因此把它与「用更小的模型直接替代大模型」混为一谈是不准确的:后者会改变输出,前者按论文说法不会。
第四,它需要额外的模型与显存。近似模型本身要占用计算与存储资源,这在大模型已经接近显存上限的部署场景中可能成为限制条件。论文并未声称该方法在任意资源约束下都适用。
第五,方法的效果与任务难度分布有关。论文的观察之一是困难任务中包含可被高效模型近似的简单子任务;如果目标任务整体上都难以被近似模型预测,可被接受的候选就会变少。