Github
Huggingface
论文提示查找解码
更新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/}
}
