一 SpecInfer 基于模型
SpecInfer([2305.09781] SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification)
SpecInfer 投机采样利用多个小型模型(SSM)快速生成“草稿”(Draft),然后由大型模型(LLM)验证并选择,从而加速文本生成。
计算过程步骤:
-
初始化:
- 输入序列: 假设我们已经有了一个输入文本序列(例如 "The cat sat on the "),需要生成接下来的文本。
- LLM: 我们有一个强大的大型语言模型 (LLM),它能给出高质量的输出,但速度较慢。
- SSM: 我们有若干个小型语言模型 (SSM),它们速度快但可能不够准确。
-
生成草稿 (Draft Generation):
- SSM 预测: 基于当前的输入序列,第一个 SSM (SSM1) 预测接下来的一个或多个 token (例如 “mat”)。
- 构建 Token 树:
- 基于
SSM1
的预测,将mat
作为一个节点添加到 token 树的根节点 (当前输入序列 "The cat sat on the ")下,形成一个分支。 - 然后,第二个 SSM (SSM2) 基于
SSM1
的输出 (“The cat sat on the mat”) 预测接下来的一个或多个 token(例如 “and” 和 “sleepy”), 形成两个新的分支。 - 第三个 SSM(SSM3) 基于以上所有分支进行继续预测,以此类推。
- 这个过程重复进行,直到达到预定的草稿长度或者遇到其他停止条件。
- 最后我们得到了一棵 Token 树,其中每个分支都代表着一个候选的 Token 序列。
- 基于
The cat sat on the|mat/ \and sleepy/ \ / \... ... ... ...
-
验证和选择 (Verification & Selection):
- LLM 验证: LLM 开始从 Token 树的根节点开始,逐个验证每个分支。
- LLM 先验证第一个分支的第一个 token “mat”,计算其概率,如果与 LLM 自己预测的概率相近,则接受 “mat” 这个 token。
- 然后,基于 “The cat sat on the mat”,LLM 会验证下一个 token,例如 “and” 和 “sleepy”。
- 如果 LLM 认为 “and” 的概率更高,则该分支会被优先考虑。
- 如果 LLM 认为 “sleepy” 概率更高,则该分支会被优先考虑。
- 分叉处理: 如果LLM验证过程中发现,某个token,比如"and",与自身预测概率相差过大,则LLM会用自己预测结果替换掉"and",并停止验证该分支。
- 确定采用的 Tokens: LLM会尽可能地从草稿树中选择更多的Tokens,直到遇到LLM认为不合适的token,则停止该轮选择。
- LLM 验证: LLM 开始从 Token 树的根节点开始,逐个验证每个分支。
-
更新输入序列:
- 将 LLM 选择的 tokens 添加到输入序列中,例如,如果LLM最终选择了 “The cat sat on the mat and”,则输入序列更新为 “The cat sat on the mat and”。
-
重复过程:
- 重复步骤 2 到 4,不断生成新的草稿并验证,直到生成满足长度要求的完整文本。
核心思路为:通过一系列小模型 SSM(Small Speculative Model)联合预测 LLM 输出,并将这些小模型的预测输出组织为 Token 树,树中每个分支表示一个候选 Token 序列。然后 LLM 使用基于树的并行解码(Tree-Based Parallel Decoding)机制来并行的验证 Token 树中所有 Token 的正确性。SpecInfer 使用 LLM 作为 Token 树验证器而非增量解码器,这显著降低了生成式 LLM 的端到端延迟,同时可以保持模型的质量。
二 REST 基于检索
REST: Retrieval-Based Speculative Decoding
REST (Retrieval-Based Speculative Decoding) 核心思想是利用一个检索模块,从一个大的文本库中找到与当前上下文相似的文本片段,并将其作为草稿生成的“指导”,然后利用一个大型模型 (Target Model) 对草稿进行验证。
计算过程步骤:
-
初始化:
- 输入序列: 假设我们有一个输入序列(例如 “The weather today is”)
- 目标模型 (Target Model): 一个大型语言模型 (LLM),用于验证和生成高质量的输出,但速度较慢。
- 检索模块 (Retrieval Module): 一个能够从大型文本库中检索相似文本片段的模块,例如利用文本向量相似度进行检索。
- 草稿长度 (Draft Length): 设定生成草稿的最大长度。
-
检索相似片段 (Retrieval):
- 检索查询: 使用当前的输入序列(例如 “The weather today is”)作为查询,送入检索模块。
- 相似文本: 检索模块从文本库中找出与查询文本最相似的文本片段,例如:“The weather today is sunny and warm.”
- 草稿生成: 从检索到的文本片段中,提取出输入序列之后的文本作为草稿,例如 “sunny and warm.”。
-
验证和选择 (Verification & Selection):
- 目标模型验证: 目标模型逐个验证草稿中的 token。
- 验证过程:
- 目标模型从草稿的第一个 token 开始验证 (例如 “sunny”)。
- 如果目标模型认为 “sunny” 合理,则接受它,并继续验证下一个 token (例如 “and”)。
- 如果目标模型认为 “sunny” 不合理,则终止验证,并用自己的预测结果替换掉该 token。
- 接受的 Tokens: 目标模型接受通过验证的 token,并添加到最终的输出序列。
-
更新输入序列:
- 将目标模型接受的 token 添加到输入序列中。
- 使用更新后的输入序列作为下一轮循环的输入。
-
重复过程:
- 重复步骤 2 到 4,直到生成满足长度要求的完整文本。
三 EAGLE 基于特征
北京大学和微软等团队提出的投机采样方案 EAGLE(Extrapolation Algorithm for Greater Language-model Efficiency,[2401.15077] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty),其是一个无损投机采样加速方案。
目标模型是原始的大模型,草稿模型是对应的小模型。
其中的 草稿模型( Draft model ) 的 Embedding 层、LM Head 以及 Sampling 都来自原始的 LLM,而 Auto-regression Head 包含一个 全连接层( FC Layer )。简单说就是利用两个词的特征信息预测下一个词的特征,(bs, seq len, 2 * hidden dim) --> (bs, seq len, hidden dim)。
具体过程描述就是: