天下武功唯快不破:投机解码
AI摘要: 本文探讨了LLM加速技术的三大类方法:模型层面、计算层面和解码过程。详细介绍了投机解码的原理及其等价性证明,通过小模型生成候选序列并由大模型验证,提高推理速度。
根据scale law,模型参数和数据量越多,效果越好。但是参数量越多带来的副作用就是大显存占用和很慢的推理速度。
一般来说,如今LLM加速技术分为三大类:模型层面,计算层面,解码过程
模型层面加速:
- 模型剪枝:删除无关紧要的权重
- 模型量化:高精度参数没有必要,换成8bit就行
- 模型蒸馏:在特定领域,只需要LLM的一部分智能,不需要全部,那就把这部分教会给一个小模型
计算层面加速:
- 模型并行:参数太多,分割到多个GPU上
- 数据并行:数据大多,分割到多个GPU上
- 混合精度计算:必要的地方用高精度,不必要的地方用低精度,封装成torch的amp装饰器
解码过程加速:
- 投机解码:用一个小而快的模型生成多个候选答案序列,由LLM评估最合理的序列,然后接着生成
- KV cache:利用self-attention的计算原理,缓存之前计算过的KV乘积
投机解码
投机解码感觉是从计算机体系结构中得到的灵感,在设计五级流水线CPU的时候就有一个加速技术为分支预测。
投机过程
- 草稿阶段:小模型快速生成前缀序列
- 验证阶段:LLM将该前缀进行推理,计算该草稿的概率,若符合预期,采样下一个token的, 得到序列。具体而言:
- 对于小模型的生成序列, LLM计算每个token的概率为, 而小模型计算每个token的概率为
- 如果 (LLM 绝对这个token不太合理),那就按照的概率拒绝这个token,从新的概率分布重新采样该token
等价证明
上面列出了投机过程,看着解释的同,但是真的合理吗?不会降低模型智能吗?通过数学证明可知,投机采样和LLM自回归采样是等价的。
-
: LLM生成的概率分布
-
: 小模型生成的概率分布
-
一个t时刻的token ()的概率为
只存在两种情况:
- LLM验证后觉得合理,也就是, 直接接受得到这个token, 则
- LLM验证后觉得不合理,也就是拒绝并重新采样得到这个token。
- 拒绝其他所有token的概率为
- LLM再采样,正好采样到这个token的概率:
- 则
故: , 证明不下去了,数学太菜了,反正大概和是属于同一个分布