提示查找解码

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

更新：这已经添加到transformers库中。请参阅此代码示例，或者只需在model.generate(...)调用中添加prompt_lookup_num_tokens=10即可。

简要总结：我们修改了推测性解码，用简单的提示中的字符串匹配来替代草稿模型，以生成候选token序列。这在基于输入的任务中带来了显著的加速（2-4倍），且不影响输出质量。该方法可用于任何解码器模型，无需更改模型或使用外部数据存储，并可用于贪婪和采样技术。

彩色token表示在单个步骤中生成了多个token。

方法

直觉：在几个LLM用例中，当你进行基于输入的生成（如摘要、文档问答、多轮对话、代码编辑）时，LLM输入（提示）和LLM输出之间存在高n-gram重叠。这可能是实体名称、短语或代码块，LLM在生成输出时直接从输入中复制。提示查找利用这种模式来加速LLM的自回归解码。

这是一个动画，用一个例子来解释（关于推测性解码本身如何工作/提供加速的信息，请参阅这篇excellent huggingface博客）：

这是"提示查找"函数：

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个token作为我们的搜索ngram
        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 and start_idx < input_length - ngram_size:
                return input_ids[0, start_idx:end_idx]

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

从实现角度来看，我们通过用这个函数替换"草稿模型"来修改推测性解码（在hf transformers中也称为辅助生成）。

该函数的输入与草稿模型相同 - 当前生成步骤的所有token（input_ids）。然后它尝试将最后几个token与提示中较早的部分匹配。如果找到，它会返回接下来k个token延续作为candidate_input_ids或候选序列。两个参数是max_ngram_size，即在提示中查找匹配时使用的最大ngram，以及num_pred_tokens，即找到匹配后要返回的候选序列长度。

实验设置

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

数据集

我们在3个数据集上进行实验，并与简单的贪婪解码作为基线进行比较。我们专注于"基于输入的"任务，在这些任务中我们预期输入和输出之间有高度重叠 - 摘要、基于上下文的问答和多轮对话。

摘要：CNN/Dailymail 100个示例
上下文问答：来自HAGRID的100个示例。我们将所有证据连接起来形成上下文，然后进行问答
多轮对话：MT-bench，所有80个示例。这并不完全是基于输入的生成，但可以了解常规对话的性能

结果

摘要和上下文问答

在摘要和上下文问答中，我们都获得了相对一致的2.4倍加速（平均）。误差条是标准差，显示根据示例不同有相当大的变化。PLD的吞吐量始终高于贪婪（或在误差范围内） - 我从未在任何示例中看到它给出比贪婪更差的吞吐量。

多轮对话

在MT-Bench上，我们在第一轮看到类似的增益，但在第零轮看到的增益要小得多。这是预料之中的 - 在第一轮中，算法只能与自己的输出匹配n-gram，因为提示相当小。然而，这种与自身输出的匹配仍然带来了可测量的增益。同样，误差条是标准差，我没有看到PLD在任何示例中给出更差的吞吐量。

MT-Bench还有提示类别。一些观察：

角色扮演的增益最差。这可能是因为没有太多n-gram可以复制，因为每次生成都是独特的。
编码在第二轮有很高的增益，因为有大量的代码复制。
在第一轮中，提取有最高的增益。这与我们的假设一致 - 在提取中肯定有n-gram复制，PLD应该有帮助。

待办事项/想法/未来工作

可能有比当前更好的字符串匹配方法，还有一些明显的改进点，例如当有多个匹配时该怎么办？最理想的延续长度是多少？
我们还没有尝试采样，尽管没有理由它不应该工作。
- 这里，一个额外需要测试的是提示查找在采样时是否会影响幻觉率，因为这人为地增加了从输入中采样精确序列的概率（这是我的同事Shwetha S提出的）
需要测试实际的FLOPs影响和权衡
还需要确定最佳超参数 - 3和10是在很少的测试基础上选择的
设计"最佳查找函数"用于解码将是一个有趣的挑战，甚至可以成为一个比赛？

如何引用

@misc{saxena2023prompt,
    title = {Prompt Lookup Decoding},
    author = {Apoorv Saxena},
    year = {2023},
    month = {November},
    url = {https://github.com/apoorvumang/prompt-lookup-decoding/}
}