提示查找解码

更新 2：此方法现已在 vLLM 中也可用，只需设置 speculative_model="[ngram]" 即可 🥳

更新：此功能已添加到 transformers 库。请参阅此处的代码示例，或直接在 model.generate(...) 调用中添加 prompt_lookup_num_tokens=10。

最小实现：请参阅演示笔记本或Colab。

简而言之：我们对推测解码进行了修改，将草案模型替换为在提示中进行简单的字符串匹配，以生成候选标记序列。这在输入相关的任务中带来了显著的速度提升（2–4倍），且不会影响输出质量。该方法可与任何解码器模型配合使用，无需更改模型或依赖外部数据存储，并且适用于贪婪和采样两种解码策略。

https://github.com/apoorvumang/prompt-lookup-decoding/assets/1957903/e908de89-ce5c-4156-8ef1-21f169dc1c8f

彩色标记表示在单步中生成了多个标记。

方法

直觉：在许多需要“基于输入生成”的大语言模型应用场景中（如摘要生成、文档问答、多轮对话、代码编辑），模型输入（提示）与输出之间存在较高的 n 元组重叠。这些可能是实体名称、短语或代码片段，模型在生成输出时会直接从输入中复制它们。提示查找利用这一模式来加速大语言模型的自回归解码。

以下是一个示例动画，用于解释其原理（有关推测解码的工作原理及其提速效果，请参阅这篇优秀的Hugging Face 博客文章）：

https://github.com/apoorvumang/prompt-lookup-decoding/assets/1957903/10c3728b-2d50-4205-a758-478e51425793

以下是“提示查找”函数：

def find_candidate_pred_tokens(input_ids, max_ngram_size=3, num_pred_tokens=10):
    input_length = input_ids.size(1)

    for ngram_size in range(max_ngram_size, 0, -1):
        # 提取最后 n 个标记作为我们的搜索 n 元组
        ngram = input_ids[0, -ngram_size:].tolist()

        # 创建大小为 ngram_size 的滑动窗口
        windows = input_ids.unfold(dimension=1, size=ngram_size, step=1)

        # 将 ngram 转换为张量以便比较
        ngram_tensor = torch.tensor(ngram, device=input_ids.device).unsqueeze(0)

        # 查找窗口与 ngram 匹配的位置
        matches = (windows == ngram_tensor).all(dim=2)

        # 获取匹配的索引
        match_indices = matches.nonzero(as_tuple=True)[1]

        # 遍历匹配索引以找到有效的延续
        for idx in match_indices:
            start_idx = idx + ngram_size
            end_idx = start_idx + num_pred_tokens
            # 确保不超过 input_ids 的长度，并避免自我匹配
            if end_idx <= input_length && start_idx < input_length - ngram_size:
                return input_ids[0, start_idx:end_idx]

    # 如果未找到匹配，则返回空张量
    return torch.tensor([], dtype=torch.long, device=input_ids.device)

从实现角度来看，我们将推测解码（也称为 Hugging Face Transformers 中的辅助生成）中的“草案模型”替换为此函数。

该函数的输入与草案模型相同——即截至当前生成步骤的所有标记 (input_ids)。随后，它会尝试将最后几个标记与提示中较早的部分进行匹配。如果找到匹配，则返回接下来的 k 个标记作为 candidate_input_ids 或候选序列。两个参数分别为 max_ngram_size，即在提示中寻找匹配时使用的最大 n 元组大小；以及 num_pred_tokens，即找到匹配后要返回的候选序列长度。

实验设置

• GPU：单块 A100 40GB
• 模型：Mistral-7B-Instruct-v0.1
• 解码类型：贪婪解码
• 超参数：最大匹配 n 元组大小 = 3，延续长度 = 10

数据集

我们在 3 个数据集上进行了实验，并以简单的贪婪解码作为基准进行对比。我们重点关注“基于输入”的任务，这类任务通常预期输入与输出之间有较高的重叠——例如摘要生成、基于上下文的问答以及多轮对话。

• 摘要生成：CNN/Dailymail 100 个示例
• 基于上下文的问答：来自 HAGRID 的 100 个示例。我们将所有证据拼接成上下文后再进行问答。
• 多轮对话：MT-bench，共 80 个示例。虽然这并不完全是基于输入的生成任务，但可以大致了解其在常规聊天场景下的性能。

结果

摘要生成与基于上下文的问答

在摘要生成和基于上下文的问答任务中，我们均获得了相对稳定的 2.4 倍速度提升（平均值）。误差线代表标准差，表明不同示例之间的差异较大。PLD 的吞吐量始终高于或与贪婪解码持平——在任何示例中我都未观察到 PLD 的吞吐量低于贪婪解码的情况。

多轮对话

在 MT-Bench 上，我们在第一轮中看到了类似的收益，但在第零轮中收益则小得多。这是意料之中——因为在第一轮中，算法只能将其自身的输出与 n 元组进行匹配，因为提示内容非常有限。不过，这种与自身输出的匹配仍然带来了可观的收益。同样，误差线代表标准差，我并未在任何示例中观察到 PLD 的吞吐量低于贪婪解码的情况。

MT-Bench 还按提示类别进行了划分。一些观察结果如下：

• 角色扮演类别的收益最差。这可能是因为可复制的 n 元组较少，每次生成的内容都较为独特。
• 编程类别的第二轮收益非常高，因为其中包含大量代码复制。
• 在第一轮中，信息抽取类别的收益最高。这与我们的假设一致——在信息抽取任务中确实存在 n 元组复制现象，因此 PLD 应该能够起到帮助作用。

待办事项/思考/未来工作

• 字符串匹配的方法可能还有更好的实现方式，目前的实现仍有明显可改进之处，例如：当存在多个匹配时该如何处理？理想的续写长度应该是多少？
• 我们尚未尝试采样方法，尽管从理论上讲它应该可以正常工作。
  • 在这里，还有一项需要测试的内容是：在采样过程中进行提示词查找是否会影响幻觉率，因为这会人为地增加从输入中采样到精确序列的概率（这一建议来自我的同事 Shwetha S）。
• 需要测试实际的 FLOPs 影响及权衡。
• 还需要确定最佳超参数——目前使用的 3 和 10 是基于非常有限的实验选择的。
• 设计用于解码的“最优查找函数”将是一项有趣的挑战，甚至可以组织一场竞赛？

如何引用

@misc{saxena2023prompt,
    title = {提示词查找解码},
    author = {Apoorv Saxena},
    year = {2023},
    month = {十一月},
    url = {https://github.com/apoorvumang/prompt-lookup-decoding/}
}

Prompt Lookup Decoding 快速上手指南

Prompt Lookup Decoding (PLD) 是一种优化的推测解码（Speculative Decoding）技术。它通过在输入提示词（Prompt）中进行简单的字符串匹配来生成候选 token 序列，从而替代传统的草稿模型（Draft Model）。该方法特别适用于摘要、文档问答、多轮对话和代码编辑等“基于输入的生成”任务，可实现 2-4 倍的推理加速，且不影响输出质量。

目前该技术已集成至 Hugging Face transformers 库及 vLLM 框架中。

环境准备

• 系统要求：Linux / macOS / Windows
• 硬件要求：支持 CUDA 的 NVIDIA GPU（推荐 A100 或同等算力显卡以获得最佳加速效果）
• 前置依赖：
  • Python >= 3.8
  • PyTorch
  • Hugging Face transformers (建议安装最新版本以支持原生参数)

安装步骤

请确保已安装基础的 PyTorch 环境。若需使用最新功能，建议升级 transformers 库。

使用 pip 安装/升级：

pip install --upgrade transformers torch accelerate

国内加速建议：如果下载速度较慢，可使用清华或阿里镜像源：

pip install --upgrade transformers torch accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

方法一：使用 Hugging Face Transformers (推荐)

自更新后，transformers 库已原生支持该功能。只需在调用 model.generate() 时添加 prompt_lookup_num_tokens 参数即可。

代码示例：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 加载模型和分词器
model_name = "mistralai/Mistral-7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# 2. 准备输入文本 (基于输入的任务效果最佳，如摘要、问答)
prompt = "Summarize the following text: Artificial intelligence is a wonderful field..."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 3. 生成文本
# 关键参数：prompt_lookup_num_tokens
# 设置该参数即启用 Prompt Lookup Decoding
# num_pred_tokens (此处为 10) 表示匹配成功后预生成的候选 token 数量
outputs = model.generate(
    **inputs,
    max_new_tokens=100,
    prompt_lookup_num_tokens=10,  # 启用 PLD
    do_sample=False               # 当前主要支持贪婪解码 (Greedy)
)

# 4. 解码输出
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)

方法二：使用 vLLM

如果你使用 vLLM 进行部署，可以通过设置 speculative_model 为 [ngram] 来启用此功能。

启动命令示例：

python -m vllm.entrypoints.api_server \
    --model mistralai/Mistral-7B-Instruct-v0.1 \
    --speculative_model "[ngram]" \
    --num_speculative_tokens 10

核心参数说明

• prompt_lookup_num_tokens (Transformers) / --num_speculative_tokens (vLLM):
  • 含义：当在 Prompt 中找到匹配的 n-gram 后，向后预取的 token 数量。
  • 默认实验值：10。可根据具体任务调整，过大会增加验证开销，过小则加速效果不明显。
• max_ngram_size (仅在使用自定义实现时需关注):
  • 含义：用于匹配的最大 n-gram 长度。
  • 默认实验值：3。

适用场景建议

• 最佳场景：摘要生成、基于文档的问答 (RAG)、多轮对话（尤其是第二轮及以后）、代码补全/编辑。这些场景中输出内容与输入内容有很高的 n-gram 重叠。
• 效果一般场景：创意写作、角色扮演（Roleplay），因为此类任务输出与输入的重叠度较低。