什么是投机解码？

一山秋叶

已于 2025-04-22 19:50:13 修改

阅读量1k

点赞数 18

分类专栏：今天又读了洋鬼子论文，请祖国放心！文章标签：自然语言处理

于 2025-02-08 23:19:36 首次发布

本文链接：https://blog.youkuaiyun.com/m0_56695799/article/details/145155843

版权

今天又读了洋鬼子论文，请祖国放心！专栏收录该内容

11 篇文章

订阅专栏

写在最前面：文章中有许多“他山之玉”，都附上了原文链接，十分感谢大佬们的文章让我学到了很多东西🙇‍，如有侵权请联系我，我立马删帖。

(这一篇的产期太长，还请大佬们斧正)

LLM的解码过程

参考文章Under the hood of Large Language Models- part2- Decoding

自回归？

现有大语言模型（LLM）大多是自回归模型（autoregressive model），生成过程依赖于模型自身的先前输出作为新的输入，也就是说每新产生一个token，就会被放入现有的token序列，然后用新的token序列来产生下一个token，直到产生终止符。

配备编码器和解码器的架构（Google的Transformer、BERT + GPT 的变种等）广泛用于源序列到目标序列的转换（seq2seq），做“翻译”任务；而LLM“自回归”的原因在于它们仅有单解码器架构，以GPT系列为代表，擅长“文本生成”任务。

单解码器架构只在一个方向上查看token，训练目标是通过预测下一个词来最小化损失函数（通常是模型预测的词与实际词之间的差异，例如交叉熵损失），简单有效，无需额外标签或人工标注，能在非常大规模的数据集上进行预训练。

概率分布？

LLM每次的输出都是从词汇表中选取某个token的概率分布，而词汇表有哪些token在预训练时便已确定。

输入n个d维的嵌入向量，经过基于自注意力机制的Transformer块后，输出形状不变，仍是(n*d)，然后经过“语言建模头”，即一个全连接层，转换成n个V维向量，V为词汇表大小。第n个V维向量 可以看成本轮中每个token所得的原始分数，再通过softmax操作将转化成每个token这次被选取的概率值：
假设你有向量，则有

选哪个/些？

如果每次都选择得分最高者作为当前token，不失为一种好办法，正所谓“贪心策略”，但也存在一些问题：
1. 针对相同的输入，每次得到的输出都是一样的，这用到分类任务或许合理，但对于“释义”、“总结”这类生成式任务就不合适了。
2. 条件太过苛刻，计算成本高。
3. 贪心算法往往会过快地收敛到局部而非全局最优。

下面选取microsoft/Phi-3-mini-4k-instruct作为实验的LLM对象，看如何解决上述问题。

Phi-3-mini-4k-instruct是一个怎样的模型？

Phi-3-Mini-4K-Instruct 是微软推出的基于 Transformer 架构的预训练语言模型，它经过了专门的调优，以便更好地理解和生成针对特定任务的指令和回复。

这个模型的主要特点是小型化和针对指令优化，它被设计成能在资源有限的环境下运行，并且能够高效地执行与任务相关的生成任务。通常使用大量的文本数据（Phi-3数据集）进行预训练，其中包括合成数据和过滤的公开网站数据

怎么跑起这个模型？

希望从开源平台Hugging Face上拉取这个预训练好的模型和对应的分词器，可以通过镜像手动下载，参考博文镜像，或者用魔法访问网址，参考博文魔法。调用模型时用gpu会比cpu快很多，驱动太久的话博文gpu驱动老旧也许有点帮助……

贪心解码（greedy decoding）

当繁琐的环境都准备好后，试着运行下列的代码

from transformers import AutoTokenizer, AutoModelForCausalLM # GPT系列、Phi系列都属于因果/自回归语言模型
from transformers import pipeline

model_name = "microsoft/Phi-3-mini-4k-instruct"
revision="main" # 固定模型版本

# model = AutoModelForCausalLM.from_pretrained(
#     model_name,
#     device_map = "cuda",
#     torch_dtype = "auto", # 自动选择数据类型，一般常见float32或float16
#     trust_remote_code = True, # 信任远程代码，加载数据模型
#     revision = revision
# )

# tokenizer = AutoTokenizer.from_pretrained(model_name, revision = revision) # 加载与该模型兼容的分词器

model = AutoModelForCausalLM.from_pretrained(model_name) # 会自动找gpu，没有的话或者驱动太老就用cpu
tokenizer = AutoTokenizer.from_pretrained(model_name)

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    return_full_text = False, # 返回的结果仅包含生成的文本，不含原始输入
    max_new_tokens = 50, # 限制生成的文本长度

    do_sample = False # 进行采样的话生成会有更多随机性和多样性

    # do_sample = True,
    # top_k = 1

    # do_sample = True,
    # top_k = 6
    
    # do_sample = False,
    # num_beams = 5
)

# print(model)

sys_prompt = "You are a helpful assistant with expertise in paraphrasing customer reviews into similar sounding reviews"
message_1 = [
    {'role': 'system', 'content': sys_prompt},
    # 系统消息：提供了模型的任务说明（即帮助改写评论）
    {'role': 'user', 'content': 'I have been using it for 6 months and still works like a charm'}
    # 用户提供的输入文本，描述了他们对某个产品的使用体验
]

import time
# 生成下一个消息，输出生成的文本
def generate_next_message(generator, message, num_runs=1):
    start_time = time.time()
    for _ in range(num_runs):
        output = generator(message) # 返回一个字典的列表，字典的键“generated_text”存储了模型生成的文本
        print(output[0]['generated_text'])
        print("####")
    end_time = time.time()
    print(f"耗时{end_time-start_time}")

generate_next_message(generator, message_1, num_runs=5) # 调用5次，每次模型回答的都是同一句话
# 如何利用 Hugging Face Transformers 库快速加载并使用预训练的语言模型进行特定任务（如文本生成或改写），如上↑

do_sample=False的代码逻辑是，每次直接将原始分数logit最高者作为下一个token（transformers/generation/utils.py下_sample()的代码）

# token selection
if do_sample:
    probs = nn.functional.softmax(next_token_scores, dim=-1)
    # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
else:
    next_tokens = torch.argmax(next_token_scores, dim=-1)

这样的结果是什么呢？就是问了模型5次，他都是一样的回答

设置do_sample=True、top_k=1时，会将所有的logits经过softmax()变成概率分布probs，交给multinomial()后由它随机采样出下一个token（概率越高者越容易被选中），耗时稍多，每次的结果都极大概率和贪心解码时的一样

ps：multinomial()的函数定义如下

def multinomial(
    input: Tensor, # 概率分布集合
    num_samples: _int, # 待抽样本数
    replacement: _bool = False, # 不允许重复抽样
    *, 
    generator: Optional[Generator] = None, # 随机数生成器
    out: Optional[Tensor] = None # 采样结果的索引存储在这
) -> Tensor:

class GenerationMixin

一个包含用于自回归文本生成的所有函数的类！

greedy decoding：num_beams=1；do_sample=False（👆）

contrastive search：penalty_alpha>0；top_k>1（👇）

multinomial sampling：num_beams=1；do_sample=True

beam-search decoding：num_beams>1；do_sample=False（👇👇）

beam-search multinomial sampling：num_beams>1；do_sample=True

diverse beam-search decoding：num_beams>1；num_beam_groups>1

constrained beam-search decoding：constraints!=None；force_words_ids!=None

assisted decoding：if `assistant_model` or `prompt_lookup_num_tokens` is passed to `.generate()`

对比搜索（contrastive search）

当我们期待模型对于同一个问题能给出不一样的答案时，会通过“惩罚”使得相似度差异最大化。

对比搜索会在生成过程中进行回滚，即访问先前生成的内容，需要启用缓存（cache）机制，不支持有状态模型（受到历史生成内容的约束，会使回滚过程变得复杂；递归生成是逐步生成一个token，将其与先前的结果进行结合，是有状态的，而这里所谓的“比对”，是在并行生成的不同的序列样本之间的）。

参考transformers/generation/utils.py下_contrastive_search()的代码，核心步骤是

1.candidate tokens recall

2.candidate re-rank by degeneration penalty

# logit变prob
next_probs = nn.functional.softmax(processed_logit_for_next_step, dim=-1)

# 根据prob选出前top_k
top_k_probs, top_k_ids = torch.topk(next_probs, dim=-1, k=top_k)

# 计算退化惩罚、重新排序（再详细解释这个代码就真的没完没了，离主线越偏越远了）
selected_idx = _ranking_fast(context_hidden, next_hidden, top_k_probs, cosine_matrix_mask, penalty_alpha, top_k)

# 从 top-k 候选中选择最优的token，并更新模型的隐藏状态
next_tokens = top_k_ids[range(len(top_k_ids)), selected_idx] # 返回形如(batch_size,)的张量，表示每个批次（序列样本个数）的最优token
next_hidden = torch.stack(torch.split(next_hidden.squeeze(dim=1), top_k))
next_hidden = next_hidden[range(batch_size), selected_idx, :]

# 通过选中的token更新模型的缓存（past_key_values），避免重复计算
next_past_key_values = selected_outputs["past_key_values"]

在生成器（generator）设置中，将参数do_sample设为True，添加参数top_k=6（越大找出来的句子越花），看看效果

（ps：在参考论文里认为直接选出logit前5，避免的对所有logit进行softmax，可以显式减少计算量，但我再花了很多时间查看源码后，还是觉得是全部softmax后选出概率前5，大家怎么看？）

（ps：top-p这个参数就是，假设你设置top-p=0.9，模型就会选累积概率超过0.9的最小词汇集，从中随机采样出下一个token）

为什么要softmax而不直接用logit呢？

logit本身不具备概率的含义，甚至可能是负数，很大或很小这样的极端值也会影响决策过程。softmax将其归一化到[0,1]，和为1，可以提高模型的稳定性和鲁棒性。指数会放大正值差距，软化负数差距，强化大值，抑制小值，让模型做出更确定的决策。

可以引入温度参数T来调节输出分布的“软硬”程度，温度T越高越趋于均匀分布，反之则会更趋向于对最可能的类别给予更高的概率。

但也不要忽略了softmax是一个昂贵的操作，浮点数除法涉及计算倒数、精度控制等步骤，会多次迭代以获取准确结果。

束搜索（beam-search decoding）

生成器设置改成

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    return_full_text=False,
    max_new_tokens=50,
    do_sample=False,
    num_beams=5
    )

看看效果

耗时更久，句子单一，它的目的在哪？为了生成最佳序列，解决“贪心”的次优解问题，通过在每一步保留多个候选序列来平衡生成质量和计算效率。过程概述如下：

1.初始化：从输入的起始符开始，生成候选的第一个词

2.迭代生成和筛选：每一时刻，模型根据当前已生成的序列继续生成下一个词，新候选序列都有一个分数，模型只保留当前所有候选序列中得分最高的前5名（下面的示意图里就会每次保留前3）

3.结束条件：达到最大生成长度或生成了结束符号，最终返回得分最高的序列

所谓“序列的分数”是什么东西？

分数指的是每个候选序列的“质量”或“概率”，具体生成过程中定义可以有很多种形式，通常情况如下：

1.对数似然概率（log likelihood）：

对于生成序列，计算该序列的所有词的条件概率的对数和

2.归一化对数似然（normalized log likelihood）：

在长文本生成任务中，为了消除长度差异的影响，通常会对分数进行归一化，即将对数似然除以序列长度T

3.惩罚因子（penalization terms）：

重复惩罚的作用是避免模型重复生成相同的词或短语，通过降低已生成词的概率来实现。还有长度惩罚避免过长

4.其他：

有时，为了引入更多的上下文信息或进行更复杂的模型调整，分数也可能结合了其他启发式方法，比如基于先验知识的贝叶斯评分或通过模型校准得到的评分

投机解码（speculative sampling）

参考论文：

Accelerating Large Language Model Decoding with Speculative Sampling

很多解释说是由一个小模型（draft model）连续猜一顺溜草稿，再由大模型（target model）并行地验证这些结果。（论文中给出的伪代码似乎有些问题，大佬们可以帮我看看下面的修改是否正确🙏）

接受“草稿”

接受小模型产生的 $\tilde{x}_i$ 的概率为 $\min \left( 1, \frac{q(\tilde{x}_i|x_1,...,x_{t+i-1})}{p(\tilde{x}_i|x_1,...,x_{t+i-1})} \right)$ ，其中 $q(\tilde{x}_i|x_1,...,x_{t+i-1})$ 和 $p(\tilde{x}_i|x_1,...,x_{t+i-1})$ 分别为大模型和小模型基于当前语境下预测下一个token为 $\tilde{x}_i$ 的概率。

拒绝“草稿”

当 $\tilde{x}_i$ 被拒绝的时候，需要从以下分布中对 $x_{t+i}$ 进行重新采样：

$x_{t+i} \sim (q(x|x_1,...,\ x_{t+i-1})-p(x|x_1,...,\ x_{t+i-1}))_+$

其中 $(.)_+$ 是一种运算： $(f(x))_+ = \frac{\max (0, f(x))}{\sum_x \max(0, f(x))}$ ，得到归一化的非负值分布。

为何正确？

即使使用了草稿模型进行初步采样，通过修正的拒绝采样方法，最终的样本分布任然是目标模型所期待的。“modified rejection sampling recovers the target distribution”怎么做到的？

每个token最终的预测结果 $X$ 为 $x$ 只有两种情况，草稿模型说是 $x$ 然后被接受了，或者说是其它东西被拒绝后重采样到了 $x$ ，即

$\mathbb{P}(X=x) \\ = \mathbb{P}(X=x, \tilde x\ accepted) + \mathbb{P}(X=x, \tilde x\ rejected) \\ = \mathbb{P}(\tilde{x}=x, \tilde x\ accepted) + \mathbb{P}(X=x, \tilde x\ rejected) \\ = \mathbb{P}(\tilde x = x) \mathbb{P}(\tilde x\ accepted|\tilde x = x) + \mathbb{P}(\tilde x\ rejected) \mathbb{P}(X=x|\tilde x\ rejected)$

在接受的情况下

$\mathbb{P}(\tilde x = x) \mathbb{P}(\tilde x\ accepted|\tilde x = x) \\ =p(x) \min \left( 1, \frac{q(x)}{p(x)} \right ) \\ =\min (p(x), q(x))$

在拒绝的情况下

$\mathbb{P}(X=x|\tilde x\ rejected) \\ =\left(q(x) - p(x) \right )_+ \\ =\frac{\max (0, q(x)-p(x))}{\sum_{x'} \max (0, q(x')-p(x'))}$

$\mathbb{P}(\tilde{x} \ rejected) \\ = 1 - \mathbb{P}(\tilde{x} \ accepted) \\ = 1 - \sum_{x'} \mathbb{P}(X = x', \tilde{x} \ accepted) \\ = 1 - \sum_{x'} \min (p(x'), q(x')) \\ = \sum_{x'} q(x') - \sum_{x'} \min (p(x'), q(x')) \\ = \sum_{x'} q(x') - \sum_{x', p(x') > q(x')} q(x') - \sum_{x', p(x') \leq q(x')} p(x') \\ = \sum_{x'} \max (0, q(x') - p(x')) \\ \text{or} = \sum_{x'} p(x') - \sum_{x', p(x') > q(x')} q(x') - \sum_{x', p(x') \leq q(x')} p(x') \\ = \sum_{x'} \max (p(x') - q(x'), 0) \\$

（ $\left(q(x) - p(x) \right )_+$ 的分母中p、q位置不打紧，分子的则严格不能调）

于是 $\mathbb{P}(\tilde x\ rejected) \mathbb{P}(X=x|\tilde x\ rejected) = \max(0, q(x')-p(x'))$

两相结合得 $\mathbb{P}(X=x) = \min (p(x), q(x)) + \max (0, q(x) - p(x)) = q(x)$

如何高效？

对于大模型q，并行计算K个logits和采样单个token的耗时相当。

在while的每一轮循环中，预测token的个数（由t指示）最多可达到K+1，最少也有1。

针对前文各种采样方法鲁棒，只需相应调整probabilities即可。

每推理一个token，明显小模型比大模型假设更快（low enough latency），在回答一些简单问题的时候，小模型的准确度也是能打的（high enough acceptance rate），在遇到难题时才调用大模型生成，就可以保证整体的又快又好。假设我们调用了n次小模型D，1次大模型T，正确生成了m个token，平均每个耗时。只要nD显著小于(m-1)T，就能实现很好的加速效果（break-even）。所以该选取怎样的草稿模型？