Speculative Decoding
直接版本
原始论文: Fast Inference from Transformers via Speculative Decoding
核心思路:
- 小模型给出若干 token
- Token 线性排列
- 大模型验证
直观:一个生成任务的每一部分难易度是不同的,有些部分相对简单,小的模型就可以生成对的结果,有的部分相对困难,需要大的模型生成。考虑到模型组 batch 的运行时间和单条数据的运行时间相差不多,如果用小的模型生成多个 token,再用大的模型一次性检查这些 token,接受其中一部分,就可以显著降低均摊开销。
算法
考虑一个大模型 ,针对输入给出概率分布 。还有一个小尺寸的模型 ,给出概率分布 。后文将两个概率分布分别简写为 。普通的 decoding 过程,就是基于 采样得到 ,具体的采样过程与 decoding 策略有关,可能是 top-1,top-k 等等。使用 代替 进行采样:
- 基于概率分布 采样得到
- 如果 ,接受
- 如果 ,以 概率接受
- 如果 被拒绝,则按照 的分布重新采样 作为结果
可以证明这个 decoding 策略给出的 token 概率分布和原始 decoding 策略是相同的,证明放在本节末尾。
基于上述考虑,使用 生成多个 token,再使用 进行验证,找到第一个被拒绝的 token,丢弃其后的全部 token,然后接受它以及之前的全部 token 作为输出。这就使得只需要调用一次大模型,就能 decode 出多个 token。算法如下:
概率分布相等的证明
给出的 被接受的概率,记为 ,是如下的概率分布的期望:
- 因此 。 当 被拒绝时,重新采样的概率 。
假设最后输出的 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,而是在所有节点中选取 个接受概率最高的节点(每个节点的接受概率仍然是路径上的概率乘积)。由于子节点的概率一定低于父节点,因此挑选出来的 个节点一定构成一棵包含根节点的子树。把这棵子树送给大模型验算,接受其中的一部分。