推测解码(Speculative decoding)

¶什么是推测解码？

推测解码（Speculative Decoding）是一种用于加速大型语言模型（LLM）推理的创新技术。它通过使用一个较小的"草稿模型"来预测下一个 token，然后用更大的"目标模型"来验证这些预测，从而显著减少推理时间。

¶技术原理

¶基本思想

推测解码的核心思想是：

使用一个小而快的草稿模型（draft model）来生成多个候选 token
使用大而准确的目标模型（target model）来并行验证这些候选 token
接受所有被验证正确的 token，拒绝第一个错误的 token 并重新采样

¶实现细节

¶详细工作流程

推测解码的具体工作流程如下：

初始预测：目标模型（LLM）预测第一个 token
草稿生成：草稿模型基于当前上下文生成多个候选 token（通常 3-5 个）
并行验证：目标模型并行验证所有草稿模型生成的候选 token
比较接受：逐个比较目标模型和草稿模型的输出：
- 如果相同：接受该 token，继续比较下一个
- 如果不同：停止接受，进入下一轮预测
下一轮循环：从未被接受的 token 位置开始新的预测循环

¶伪代码示例

def speculative_decoding(prompt, draft_model, target_model, max_tokens, k=5):
    generated = []

    while len(generated) < max_tokens:
        # 目标模型预测第一个token（如果需要）
        if not generated:
            next_token = target_model.sample(prompt)
            generated.append(next_token)
            continue

        # 草稿模型生成k个候选token
        current_context = prompt + generated
        draft_tokens = draft_model.generate(current_context, num_tokens=k)

        # 目标模型并行验证所有候选token
        target_probs = []
        for i in range(len(draft_tokens)):
            verify_context = current_context + draft_tokens[:i]
            next_token_probs = target_model.predict_next_token(verify_context)
            target_probs.append(next_token_probs)

        # 逐个比较并接受token
        accepted_tokens = []
        for i, draft_token in enumerate(draft_tokens):
            target_prob = target_probs[i][draft_token]
            draft_prob = draft_model.get_probability(draft_token, current_context + draft_tokens[:i])

            accept_prob = min(1, target_prob / draft_prob) if draft_prob > 0 else 0

            if random.random() < accept_prob:
                accepted_tokens.append(draft_token)
            else:
                break  # 第一个不匹配就停止

        generated.extend(accepted_tokens)

        # 如果全部接受，继续下一轮；否则从目标模型重新采样
        if len(accepted_tokens) == k:
            continue
        else:
            resample_context = current_context + accepted_tokens if accepted_tokens else current_context
            next_token = target_model.sample(resample_context)
            generated.append(next_token)

    return generated