推理加速：投机采样经典方法

news2025/2/27 17:20:23

一 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 选择的 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）。