分享自:

基于推测解码的变压器快速推理

期刊:Proceedings of the 40th International Conference on Machine Learning, PMLR

本研究由 Yaniv Leviathan、Matan Kalman 和 Yossi Matias 完成,三位作者均来自 Google Research。该研究以论文形式发表,题目为 “Fast Inference from Transformers via Speculative Decoding”。本文收录于 2023 年第 40 届国际机器学习会议(ICML)的会议论文集(Proceedings of Machine Learning Research, PMLR)中。

学术背景

本研究的核心科学领域属于机器学习和人工智能,具体聚焦于大型自回归模型(如 Transformer)的推理加速。近年来,以 GPT-3、PaLM、LaMDA 等为代表的大型 Transformer 模型在文本和图像生成等任务上展现出卓越的能力。然而,这些模型的推理过程存在显著瓶颈:生成 K 个标记(tokens)需要顺序运行模型 K 次,这种串行特性导致了极高的延迟,严重限制了其在实时交互场景中的应用。

为了提升推理效率,学术界已探索多种方法。一些方法旨在同等优化所有输入的推理成本,如知识蒸馏、量化和模型架构修改;另一类自适应计算(Adaptive Computation)方法则观察到并非所有推理步骤都同等“困难”,因此尝试为“简单”步骤分配更少的计算资源。然而,这些方法通常需要修改模型架构、调整训练流程并重新训练模型,且往往无法保证与原始目标模型(Target Model)的输出分布完全一致。

本研究基于两个关键观察:(1)复杂的语言建模任务中常包含可由更高效的模型良好近似的较易子任务;(2)通过利用推测性执行(Speculative Execution)和一种新颖的采样方法,可以借助近似模型的输出来并行运行目标模型,从而在不改变输出分布的前提下加速解码。本研究旨在提出一种无需重新训练或更改模型架构,即可加速现有现成(off-the-shelf)模型推理的通用方法。

详细工作流程

本研究提出了一种名为“推测性解码”(Speculative Decoding)的新算法。其核心思想是使用一个更小、更快但功能近似的“近似模型”(Approximation Model, M_q)来推测(Speculate)生成一段候选标记序列(即“猜测”),然后由更大、更精确但更慢的“目标模型”(Target Model, M_p)并行地验证这些猜测。通过一种名为“推测性采样”(Speculative Sampling)的数学方法,系统可以接受大部分正确的猜测,并仅在必要时由目标模型进行修正,从而在单次并行计算中可能生成多个标记,显著减少目标模型的串行调用次数。

具体的工作流程(对应于论文中的算法 1: SpeculativeDecodingStep)包含以下几个关键步骤,该流程以一个给定的前缀(prefix)和预设的猜测步数 γ 作为输入:

步骤一:使用近似模型生成猜测序列。 首先,利用近似模型 M_q 以自回归的方式生成 γ 个连续的标记作为候选序列。对于 i 从 1 到 γ,算法执行:基于当前前缀(初始为给定前缀,后续为前缀加上已生成的猜测标记),运行 M_q 得到下一个标记的分布 q_i(x),然后从该分布中采样一个标记 x_i。这样就得到了一个长度为 γ 的候选标记序列 [x_1, ..., x_γ]

步骤二:并行运行目标模型进行评估。 在生成猜测序列的同时或之后,算法并行地运行目标模型 M_p 来评估这些猜测。具体而言,它需要计算: - p_1(x): 给定原始前缀,M_p 输出的下一个标记的分布。 - p_2(x): 给定原始前缀加上 x_1,Mp 输出的下一个标记的分布。 - … - p{γ+1}(x): 给定原始前缀加上完整的 γ 个猜测标记,M_p 输出的下一个标记的分布。 这一步是并行的关键,它将原本需要 γ+1 次串行计算的模型前向传播合并为一次并行计算(假设有足够的计算资源,如多个计算单元或能够批量处理)。

步骤三:确定接受多少个猜测(计算接受数 n)。 这是推测性采样的核心。对于每一个猜测位置 i(1 ≤ i ≤ γ),算法生成一个随机数 r_i ~ U(0,1),并将其与比值 p_i(x_i) / q_i(x_i) 进行比较。这里 p_i(x_i) 是目标模型认为应生成猜测 x_i 的概率,q_i(x_i) 是近似模型生成该猜测的概率。如果 r_i ≤ p_i(x_i)/q_i(x_i),则该猜测被“接受”,否则被“拒绝”。接受数 n 定义为从第一个被拒绝的猜测位置减 1,如果所有猜测都被接受,则 n = γ。这个判定规则(即当猜测概率不高估真实概率时必然接受,当高估时以一定概率拒绝)在数学上被证明能确保最终输出序列的分布与完全由目标模型 M_p 自回归生成的分布完全相同。

步骤四:调整分布并采样最终标记。 如果 n < γ(即至少有一个猜测被拒绝),则需要从第 n+1 个位置的调整后分布 p’(x) 中采样一个修正标记 t。调整规则为:p’(x) = norm(max(0, p{n+1}(x) - q{n+1}(x))),即从目标模型的分布 p{n+1}(x) 中减去近似模型分布 q{n+1}(x) 中高估的部分,然后重新归一化(norm)。如果 n = γ(所有猜测都被接受),则直接从 p_{γ+1}(x) 中采样一个额外的标记 t。

步骤五:输出结果并更新前缀。 最终,算法输出本次迭代生成的新标记序列:前 n 个被接受的猜测标记 [x_1, ..., x_n],加上最后采样的那个标记 t(可能是修正标记或额外标记)。这个新序列被附加到原始前缀后,作为下一次解码迭代的输入。

本研究涉及的研究对象为多个预训练的大规模自回归模型及其更小的近似版本。实验部分涵盖了不同任务和模型规模: 1. 无条件语言生成:使用在 LM1B 语料库上训练的 97M 参数 GPT 风格解码器作为 M_p,以及一个 6M 参数的小型同类模型作为 M_q。 2. 机器翻译和摘要:使用 T5-xxl(11B 参数)作为 M_p,并测试了 T5-small(77M)、T5-base(250M)、T5-large(800M)作为 M_q。任务为英德翻译(WMT EnDe)和新闻摘要(CNN/DM)。 3. 对话任务:使用 LaMDA(137B 参数)作为 M_p,并测试 LaMDA 的较小版本(100M,2B,8B)作为 M_q。

研究没有开发新的模型,而是完全使用现有的、公开可获取的预训练模型检查点(checkpoint)。核心的创新在于算法层面,即推测性解码算法和推测性采样方法。数据分析和实验验证完全基于对这些现成模型应用新算法后的性能测量,包括计算接受率 α、测量实际加速比(walltime improvement)以及验证输出分布的保真度。

主要结果

本研究通过理论分析和实验验证,全面评估了推测性解码方法的性能。

1. 理论分析结果: 论文推导了衡量近似模型性能的关键指标——接受率 α,定义为给定前缀下,从 M_q 采样的标记被推测性采样接受的概率期望。理论证明得出 α = E[1 - D_LK(p, q)] = E[Σ min(p(x), q(x))],其中 D_LK 是 p 和 q 分布之间的一种散度。α 值越高,说明两个模型的输出分布越接近。 基于 α 和另一个关键参数——成本系数 c(运行一次 M_q 与运行一次 M_p 的时间比),论文推导出了期望的加速比公式:加速因子 = (1 - α^{γ+1}) / ((1-α)(γc + 1))。这表明加速效果取决于:近似模型的质量(α)、其相对速度(c)以及猜测步数(γ)。分析指出,只要 α > c,就存在某个 γ 值能带来加速,并且至少能获得 (1+α)/(1+c) 的加速。论文还提供了如何根据 α 和 c 选择最优 γ 的指导(见图3)。

2. 实证结果: 接受率(α)测量: 表3展示了在不同任务、不同规模的目标模型和近似模型组合下测得的 α 值。关键发现包括: - 小两个数量级的近似模型(如 77M vs 11B)仍能产生 0.5-0.9 的 α 值,证明“易子任务”普遍存在。 - 即使使用简单的 n-gram 模型(如二元模型)作为 M_q,也能获得非零的 α(例如在英德翻译任务中 α ≈ 0.2),这意味着即使是极简单的启发式方法也能带来一定加速。 - 在相同的模型配对下,使用贪婪解码(temperature=0)通常比随机采样(temperature=1)获得更高的 α 值。

实际加速比测量: 研究在 T5-xxl(11B)模型上实现了推测性解码,并与高度优化的 T5x 实现进行了实际延迟对比(表2)。主要结果包括: - 在英德翻译任务上,使用 T5-small(77M)作为 M_q,γ=7,在随机采样(temp=1)和贪婪解码(temp=0)下分别获得了 2.6倍3.4倍 的加速。 - 在 CNN/DM 摘要任务上,使用相同的配置,分别获得了 2.3倍3.1倍 的加速。 - 加速效果与理论预测基本吻合(见表4对比),微小差异源于实现优化和“β值独立同分布”这一简化假设。

输出保真度验证: 研究强调了推测性采样方法的数学严谨性,确保了无论近似模型 M_q 如何选择,最终输出的标记序列分布与仅使用目标模型 M_p 进行标准自回归解码所得到的分布完全相同。这是该方法区别于大多数现有自适应计算方法的根本优势。

这些结果逻辑严密地支持了研究的核心主张:推测性解码是一种有效、实用且无损的推理加速方法。α 值的测量证明了使用小模型进行推测的可行性;理论公式为参数选择和性能预测提供了框架;而最终的实验数据则直接验证了该算法在实际硬件上能带来显著的延迟降低,且输出质量不变。

结论与意义

本研究提出并验证了“推测性解码”,一种能够在不改变输出分布、无需修改模型架构或重新训练的前提下,显著加速大型自回归 Transformer 模型推理的新方法。

其科学价值在于: 1. 理论创新:将处理器中的“推测性执行”思想创造性地推广到随机计算(Stochastic Setting)领域,并提出了“推测性采样”这一新颖的数学方法,保证了采样过程的无偏性。 2. 方法论创新:提供了一种通用的、即插即用的解码加速框架。该框架独立于具体的模型架构和训练过程,可应用于任何自回归生成模型。

其应用价值巨大: 1. 实用性:可以直接应用于现有的、已部署的大型语言模型(如 T5、GPT、LaMDA 等),立即获得 2-3 倍的推理速度提升,这对于降低服务成本、改善用户体验(如聊天机器人响应速度)具有重要意义。 2. 部署简便:无需训练新模型或修改现有模型,只需在推理服务中集成一个更小的近似模型并实现所述算法即可。 3. 资源利用优化:特别适用于计算瓶颈在于内存带宽而非算术运算的场景(许多大模型推理的现状),该方法通过增加计算并发性来充分利用空闲的计算单元。

研究亮点

  1. 核心创新点:提出了“推测性采样”算法,这是首个能够在保证输出分布严格不变的前提下,通过并行化加速随机自回归解码的方法。
  2. 显著性能提升:在实际的大规模模型(11B 参数)上实现了 2-3 倍的端到端延迟降低,且经过严格验证输出无误。
  3. 强大的通用性与简易性:方法完全与模型无关,支持各种采样策略(贪婪、top-k、核采样等),并能兼容非自回归模型、n-gram 模型等多种类型的近似模型,甚至随机猜测也能带来微弱加速。
  4. 扎实的理论分析:提供了完整的数学推导,包括接受率 α 的计算、加速比的理论公式、最优猜测步数 γ 的选择策略等,为理解和应用该方法奠定了坚实基础。
  5. 广泛的实验验证:在多种任务(生成、翻译、摘要、对话)、多种模型架构(仅解码器、编码器-解码器)和不同规模(从数百万到上千亿参数)上进行了验证,证明了方法的普适性。

其他有价值内容

论文还探讨了该方法的局限性与未来方向: - 局限性:加速是通过增加算术运算的并发性(从而增加总运算量)来换取的,因此在计算资源(而非内存带宽)受限的场景中可能不适用。 - 与束搜索(Beam Search)的结合:附录中简要探讨了将推测性解码应用于束搜索的可能性,虽然会带来一些性能惩罚,但思路是可行的,留待未来研究。 - 宽松度(Lenience)参数:附录提出,如果允许输出分布有轻微、可控的偏离(例如,保证任何标记的输出概率不超过其真实概率的 2 倍),则可以引入一个“宽松度”参数来进一步提高接受率 α,从而获得更大的加速潜力(例如理论加速比可达 5 倍)。这为在输出质量与推理速度之间进行权衡提供了新的维度。 - 未来工作:包括探索分层推测(用更快的模型加速 M_q)、动态调整 γ、为特定任务定制训练近似模型,以及将该“随机推测性执行”思想应用于强化学习、物理模拟等其他领域。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com