Fast Inference from Transformers via Speculative Decoding

type

status

date

slug

summary

摘要

Inference from large autoregressive models like Transformers is slow - decoding K tokens takes K serial runs of the model. In this work we introduce speculative decoding - an algorithm to sample from autoregressive models faster without any changes to the outputs, by computing several tokens in parallel. At the heart of our approach lie the observations that (1) hard language-modeling tasks often include easier subtasks that can be approximated well by more efficient models, and (2) using speculative execution and a novel sampling method, we can make exact decoding from the large models faster, by running them in parallel on the outputs of the approximation models, potentially generating several tokens concurrently, and without changing the distribution. Our method can accelerate existing off-the-shelf models without retraining or architecture changes. We demonstrate it on T5-XXL and show a 2X-3X acceleration compared to the standard T5X implementation, with identical outputs.

像 Transformers 这样的大型自回归模型的推断速度很慢--解码 K 个 token 需要对模型进行 K 次串行运行。在这项工作中，我们引入了投机解码--一种通过并行计算多个标记，在不改变输出的情况下更快地从自回归模型中采样的算法。我们的方法的核心在于观察到：(1) 困难的语言建模任务通常包括更容易的子任务，这些子任务可以用更高效的模型很好地近似；(2) 使用投机执行和一种新颖的采样方法，我们可以通过在近似模型的输出上并行运行大型模型，在不改变分布的情况下同时生成多个标记，从而更快地从大型模型中进行精确解码。我们的方法可以加速现有的现成模型，而无需重新训练或改变架构。我们在 T5-XXL 上进行了演示，结果表明，与标准 T5X 实现相比，在输出相同的情况下，我们的方法可加速 2 倍至 3 倍。

一简介

关键观察：

有些推理步骤是“困难的”，有些是“简单的”

大模型推理性能瓶颈是内存带宽和通信的而不是算术运算的，意味着有额外的计算资源可用。因此，可以增加一些并发行为补充计算量。

投机执行通常是处理器中的一种优化技术，目的是去验证一个任务是否确实需要并行执行，任务回报是增加并发性。经典例子：分支预测。

应用在自回归模型上：从近似模型中采样作为推测前缀，用投机采样的方法最大化任务被接受的可能性，同时保证和原始模型保持同样的结果分布。

特点：

不需要训练新模型（区别于 Blockwise parallel decoding）

不改变输出

贡献：

将推测执行推广到随即设置，采用新方法——投机采样

采用投机解码加速自回归模型解码，同时不改变模型架构、训练流程、输出分布

二 Speculative Decoding

2.1 概述

是目标模型（原本的大模型），是模型对前缀的输出分布，下文用表示。

是更高效的近似模型，输出分布是。

核心思想：

用生成个结果

用并行评估的所有猜测和其各自的概率，接受所有可能导致相同分布的猜测

从调整后的分布中采样一个额外 token 来修正第一个被拒绝的 token，或者他们全部被接受则添加一个额外 token

这样，每次并行运行就可以产生至少 1 个，至多个 token。

2.2 标准采样

虽然有许多采样方法和参数，例如 top-k、top-p，并且流行的实现通常在 logits 级别上对它们进行不同的处理，但它们都可以转换为标准采样调整后的概率分布。

例如，采样相当于将分布的非最大元素清零并标准化。

2.3 投机采样

为了采样，我们用代替。如果就接受，如果就以的概率拒绝，并从修正后的分布中重新采样。

举个例子：

在上的运行结果为，采样

在和上并行运行，得到和

如果被拒绝，则忽略，并从修正的分布中重新采样；如果被接受，则保留和

略

三分析（略）

3.1 token 生成数量

💡

定义 3.1：给定前缀，接受率为投机采样接受的概率

是与近似程度的自然度量。

如果简化假设是独立同分布的，记，则算法 1 单次运行产生的 token 数量是一个有上限的几何变量，其成功概率为，上限为。

Token 生成数量的期望值满足公式（1）。如图 2

The expected number of tokens generated by Algorithm 1 as a function of α for various values of γ

3.2 计算

给定前缀和两个模型，计算。定义一个自然散度：

💡

定义 3.2：，其中

💡

引理 3.3： 证明：

💡

：

从引理 3.3 可以得到推论 3.4：

💡

推论 3.4： 是一个在的对称散度

💡

定理 3.5： 证明：

💡

存疑

💡

推论 3.6： 转化成标准采样？

3.3 ~ 3.5

时间提升

算术运算次数

选择

摘要

一简介