Contents

Speculative Decoding

直接版本

原始论文: Fast Inference from Transformers via Speculative Decoding

核心思路:

  • 小模型给出若干 token
  • Token 线性排列
  • 大模型验证

直观:一个生成任务的每一部分难易度是不同的,有些部分相对简单,小的模型就可以生成对的结果,有的部分相对困难,需要大的模型生成。考虑到模型组 batch 的运行时间和单条数据的运行时间相差不多,如果用小的模型生成多个 token,再用大的模型一次性检查这些 token,接受其中一部分,就可以显著降低均摊开销。

算法

考虑一个大模型 ,针对输入给出概率分布 。还有一个小尺寸的模型 ,给出概率分布 。后文将两个概率分布分别简写为 。普通的 decoding 过程,就是基于 采样得到 ,具体的采样过程与 decoding 策略有关,可能是 top-1,top-k 等等。使用 代替 进行采样:

  1. 基于概率分布 采样得到
  2. 如果 ,接受
  3. 如果 ,以 概率接受
  4. 如果 被拒绝,则按照 的分布重新采样 作为结果

可以证明这个 decoding 策略给出的 token 概率分布和原始 decoding 策略是相同的,证明放在本节末尾。

基于上述考虑,使用 生成多个 token,再使用 进行验证,找到第一个被拒绝的 token,丢弃其后的全部 token,然后接受它以及之前的全部 token 作为输出。这就使得只需要调用一次大模型,就能 decode 出多个 token。算法如下:

/images/spec_decode_algo.png

概率分布相等的证明

给出的 被接受的概率,记为 ,是如下的概率分布的期望:

  • 因此 。 当 被拒绝时,重新采样的概率

假设最后输出的 token 是 ,输出该 token 的概率是 其中 带入得到 这就是在说,最终输出的 token 概率分布,和原始模型给出的概率分布 完全等同。

EAGLE

原始论文: EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees

核心思路:

  • 小模型不是给出 token 而是给出高维向量,然后用大模型的最后一层把向量映射到概率分布。也就是说,小模型在预测大模型倒数第二层的输出,而非最后一层的输出。
  • 小模型不是给出一个线性的预测,而是给出一棵树,大模型验证后接受树上的一条路径。
  • (EAGLE-2)动态的树结构,根据小模型的打分决定树的结构,把预测的 budget 放在更合适的地方。

记号

token:,token embedding:,feature(倒数第二层输出的隐藏向量),最后一层输出的概率分布

原始的 decoding 过程可以如下描述:

Drafting Phase

Draft model 使用 来预测 ,进而给出 。之所以 token 和 feature 错开了一位,是因为 attention 机制给不出第一个 token 之前的 feature。用符号语言描述就是

训练这个 draft model 使用一个复合的 loss 函数: 其中 是一个混合系数,取 表示 feature 拟合误差, 表示 token 概率分布分布误差: 在训练时为了避免错误累计,要给训练集的 feature 加上一个均匀随机的噪音

我的看法: 似乎本质上是一样的。但是需要二者结合来用,可能说明了在 feature 层面的误差经过 LM_Head 会显著放大。

在 drafting 阶段,算法并不是保持一个线性的猜测,而是给出一个树状结构。每一轮 draft model forward,都从树上挑选一些叶子节点进行扩展。可以是贪心地挑选概率最高的 个节点,也可以是依照某个分布采样出 个节点。最终多轮 forward 后,会形成一个不太平衡的树状结构。这棵树的结构是相对固定的,每一层有多少个节点多少分支都是确定的。

Verification Phase

draft 形成的树上的一条路径上是有因果关系的,树上的兄弟节点之间没有直接因果关系。因此在把树上的全部节点送给大模型进行验证时,需要设置特殊的 causal mask,才能保证结果正确。

EAGLE-2 的改进

直观来说,生成任务的每一个部分难度是不一样的。树上的某些节点,它处于容易的上下文之中,它对应的概率分布集中于一个点,那么这个节点就不该分出两个或以上的分支,在那些明显低概率的地方浪费验算阶段的算力。而某些节点的概率比较均匀地散开在几个候选项上,那么这些节点就可以多分出几个分支。因此,根据小模型打分,动态地决定这个树的结构,可以节省算力。

实验发现,小模型给出的 token 概率分布,和对应 token 在验证阶段被接受的概率很接近。因此可以直接使用小模型给出的 token 概率,来估算它的真实接受概率。基于这个估算的概率分数,可以动态地控制 draft tree 的结构。draft tree 按照如下规则生成:

  • 扩展阶段。选取当前层接受概率最大的 个节点进行扩展。一个 token 被接受,当且仅当它以及前面的所有 token 被接受。因此这个接受概率要按照到根节点的路径求乘积。
  • 重排阶段。并不是直接把上一阶段扩展过的路径拿来当作 draft,而是在所有节点中选取 个接受概率最高的节点(每个节点的接受概率仍然是路径上的概率乘积)。由于子节点的概率一定低于父节点,因此挑选出来的 个节点一定构成一棵包含根节点的子树。把这棵子树送给大模型验算,接受其中的一部分。