大模型系列——投机解码:Prompt Lookup Decoding代码解读

官方代码见:GitHub - apoorvumang/prompt-lookup-decoding

UPDATE 2 : This method is now available in vLLM as well by setting speculative_model="[ngram]" 🥳

UPDATE : This has been added to the transformers library. Please see this for a code example, or simply add prompt_lookup_num_tokens=10 to your model.generate(...) call.

TLDR : We modify speculative decoding where we replace the draft model with simple string matching in the prompt to generate candidate token sequences. This results in significant speedups (2x-4x) in input-grounded tasks, with no effect on output quality. This method can be used with any decoder model without model changes or external datastore, and with both greedy and sampling techniques.

Intuition : In several LLM use cases where you're doing input grounded generation (summarization, document QA, multi-turn chat, code editing), there is high n-gram overlap between LLM input (prompt) and LLM output. This could be entity names, phrases, or code chunks that the LLM directly copies from the input while generating the output. Prompt lookup exploits this pattern to speed up autoregressive decoding in LLMs.

python 复制代码
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):
        # Extract the last n tokens as our search ngram
        ngram = input_ids[0, -ngram_size:].tolist()

        # Create sliding windows of size ngram_size
        windows = input_ids.unfold(dimension=1, size=ngram_size, step=1)

        # Convert ngram to a tensor for comparison
        ngram_tensor = torch.tensor(ngram, device=input_ids.device).unsqueeze(0)

        # Find where the windows match the ngram
        matches = (windows == ngram_tensor).all(dim=2)

        # Get the indices of matches
        match_indices = matches.nonzero(as_tuple=True)[1]

        # Iterate through match indices to find a valid continuation
        for idx in match_indices:
            start_idx = idx + ngram_size
            end_idx = start_idx + num_pred_tokens
            # Ensure we don't go beyond the length of input_ids and avoid self-match
            if end_idx <= input_length and start_idx < input_length - ngram_size:
                return input_ids[0, start_idx:end_idx]

    # If no match is found, return an empty tensor
    return torch.tensor([], dtype=torch.long, device=input_ids.device)

ODOs/Thoughts/Future work

  • There's probably better ways to do stringmatching than the current one, and there are several obvious things to improve eg. what to do when there are multiple matches? Whats the ideal length of continuation?
  • We haven't yet tried sampling, although there's no reason it shouldn't work.
    • Here, one additional thing to test would be whether prompt lookup while sampling can affect hallucination rates, since this artifically increases probability of sampling exact sequences from input (this was suggest by my colleague Shwetha S)
  • Testing actual FLOPs impact and tradeoffs is needed
  • Also need to figure out best hyperparams - 3 and 10 were chosen on very little testing
  • It would be an interesting challenge to design the "best lookup function" for decoding, could even be a competition?

这个方法可能还是有问题的,正如坐着所说,可能存在幻觉,不一定ngram匹配上的就能加速

相关推荐
CV视觉15 小时前
智能体综述:探索基于大型语言模型的智能体:定义、方法与前景
人工智能·语言模型·chatgpt·stable diffusion·prompt·aigc·agi
小关会打代码2 天前
LangChain最详细教程之Model I/O(二)Prompt Template
langchain·prompt
Tiandaren3 天前
自用提示词01 || Prompt Engineering || 学习路线大纲 || 作用:通过启发式的问题来带动学习
人工智能·pytorch·深度学习·nlp·prompt·1024程序员节
梵得儿SHI4 天前
大型语言模型基础之 Prompt Engineering:打造稳定输出 JSON 格式的天气预报 Prompt
人工智能·语言模型·prompt·提示词工程·结构化输出·engineering·ai交互
梵得儿SHI4 天前
Prompt Engineering 关键技能:精准掌控 LLM 输出的格式、内容与风格
大模型·llm·prompt·格式控制·内容到风格·内容控制·风格控制
L_1421906875 天前
论文泛读:DYNAPROMPT: DYNAMIC TEST-TIME PROMPT TUNING(动态测试时调优)
prompt
文火冰糖的硅基工坊5 天前
[人工智能-大模型-21]:“AI 编程工作流”模板(含 prompt 示例)
人工智能·科技·学习·大模型·prompt
亿牛云爬虫专家5 天前
用LLM自动化生成解析器:从Prompt到Parser的工程化落地
自动化·llm·prompt·xpath·爬虫代理·网页爬虫·解析器
理不为5 天前
提示词 prompt 快速上手
人工智能·prompt
猫头虎6 天前
openAI发布的AI浏览器:什么是Atlas?(含 ChatGPT 浏览功能)macOS 离线下载安装Atlas完整教程
人工智能·macos·chatgpt·langchain·prompt·aigc·agi