天下武功唯快不破:投机解码

·1562·4 分钟·
AI摘要: 本文探讨了LLM加速技术的三大类方法:模型层面、计算层面和解码过程。详细介绍了投机解码的原理及其等价性证明,通过小模型生成候选序列并由大模型验证,提高推理速度。

根据scale law,模型参数和数据量越多,效果越好。但是参数量越多带来的副作用就是大显存占用和很慢的推理速度。

一般来说,如今LLM加速技术分为三大类:模型层面,计算层面,解码过程

模型层面加速:

  1. 模型剪枝:删除无关紧要的权重
  2. 模型量化:高精度参数没有必要,换成8bit就行
  3. 模型蒸馏:在特定领域,只需要LLM的一部分智能,不需要全部,那就把这部分教会给一个小模型

计算层面加速:

  1. 模型并行:参数太多,分割到多个GPU上
  2. 数据并行:数据大多,分割到多个GPU上
  3. 混合精度计算:必要的地方用高精度,不必要的地方用低精度,封装成torch的amp装饰器

解码过程加速:

  1. 投机解码:用一个小而快的模型生成多个候选答案序列,由LLM评估最合理的序列,然后接着生成
  2. KV cache:利用self-attention的计算原理,缓存之前计算过的KV乘积

投机解码

投机解码感觉是从计算机体系结构中得到的灵感,在设计五级流水线CPU的时候就有一个加速技术为分支预测。

投机过程

  1. 草稿阶段:小模型快速生成前缀序列x1,x2,x3,...,xtx_1, x_2, x_3, ..., x_t
  2. 验证阶段:LLM将该前缀进行推理,计算该草稿的概率,若符合预期,采样下一个token的xt+1x_{t + 1}, 得到序列x1,x2,...xt,xt+1x_1, x_2, ... x_{t}, x_{t + 1}。具体而言:
    1. 对于小模型的生成序列x1...tx_{1 ... t}, LLM计算每个token的概率为p(xt)p(x_t), 而小模型计算每个token的概率为q(xt)q(x_t)
    2. 如果p(xt)<q(xt)p(x_t) < q(x_t) (LLM 绝对这个token不太合理),那就按照1p(xt)q(xt)1 - \frac{p(x_t)}{q(x_t)}的概率拒绝这个token,从新的概率分布p(x)=norm(max(0,p(xt)q(xt)))p^`(x) = norm(max(0, p(x_t) - q(x_t)))重新采样该token

等价证明

上面列出了投机过程,看着解释的同,但是真的合理吗?不会降低模型智能吗?通过数学证明可知,投机采样和LLM自回归采样是等价的。

  • p(x)p(x): LLM生成的概率分布

  • q(x)q(x): 小模型生成的概率分布

  • 一个t时刻的token (xtx_t)的概率为β(x)\beta(x)

β(xt)\beta(x_t) 只存在两种情况:

  1. LLM验证后觉得合理,也就是p(x)>=q(x)p(x) >= q(x), 直接接受得到这个token, 则 β(xt)=q(xt)=min{p(xt),q(xt)}\beta(x_t) = q(x_t) = min\{p(x_t), q(x_t)\}
  2. LLM验证后觉得不合理,也就是p(x)<q(x)p(x) < q(x)拒绝并重新采样得到这个token。
    1. 拒绝其他所有token的概率为1xt!=xtmin{p(xt),q(xt)}1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\}
    2. LLM再采样,正好采样到这个token的概率:p(xt)p(x_t)
    3. β(xt)=p(xt)×(1xt!=xtmin{p(xt),q(xt)})\beta(x_t) = p(x_t) \times (1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\})

故:β(xt)=min{p(xt),q(xt)}+p(xt)×(1xt!=xtmin{p(xt),q(xt)})\beta(x_t) = min\{p(x_t), q(x_t)\} + p(x_t) \times (1 - \sum_{x^`_t != x_t} min\{p(x_{t^`} ), q(x_{t^`})\}) , 证明不下去了,数学太菜了,反正大概和p(xt)p(x_t)是属于同一个分布