beam search、top-k sampling、nucleus sampling、temperature sampling和联合采样

taoqick

已于 2024-12-03 12:00:49 修改

阅读量2.6k

点赞数 11

分类专栏：机器学习算法文章标签：深度学习

于 2022-04-01 16:00:14 首次发布

原文链接：https://zhuanlan.zhihu.com/p/82829880

版权

算法同时被 2 个专栏收录

475 篇文章

订阅专栏

机器学习

74 篇文章

订阅专栏

这几种解码策略在hugging face的GenerationMixin(transformers/generation/utils.py)中均有所实现，在hugging face上的生成式模型都要继承GenerationMixin，以beamsearch为例，下面self就是继承的子类提供的根据w_{0..i-1}给w_{i}打分的language model，这个language model里当然要实现例如kv_cache等策略：

            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

            outputs = self(
                **model_inputs,
                return_dict=True,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
            )

假设一个搜索任务

假设现在有一个简化版的中文翻译英文任务，输入和输出如下，为了方便描述搜索算法，限制输出词典只有{"I", "H", "U"} 这3个候选词，限制1个时间步长翻译1个汉字，1个汉字对应1个英文单词，这里总共3个汉字，所以只有3个时间步长。

中文输入："我" "恨" "你"
英文输出："I" "H" "U"

目标：得到最优的翻译序列 I-H-U

exhaustive search（穷举搜索）

最直观的方法就是穷举所有可能的输出序列，3个时间步长，每个步长3种选择，共计种排列组合。

I-I-I
I-I-H
I-I-U
I-H-I
I-H-H
I-H-U
I-U-I
I-U-H
I-U-U

H-I-I
H-I-H
H-I-U
H-H-I
H-H-H
H-H-U
H-U-I
H-U-H
H-U-U

U-I-I
U-I-H
U-I-U
U-H-I
U-H-H
U-H-U
U-U-I
U-U-H
U-U-U

从所有的排列组合中找到输出条件概率最大的序列。穷举搜索能保证全局最优，但计算复杂度太高，当输出词典稍微大一点根本无法使用。

greedy search（贪心搜索）

贪心算法在翻译每个字的时候，直接选择条件概率最大的候选值作为当前最优。如下图所以，

第1个时间步长：首先翻译"我"，发现候选"I"的条件概率最大为0.6，所以第一个步长直接翻译成了"I"。
第2个时间步长：翻译"我恨"，发现II概率0.2，IH概率0.7，IU概率0.1，所以选择IH作为当前步长最优翻译结果。
第3个时间步长：翻译"我恨你"，发现IHI概率0.05，IHH概率0.05，IHU概率0.9，所以选择IHU作为最终的翻译结果。

PS：图中的概率如何得来的？不同的模型有不同的算法，我自己随便填的。

greedy search

贪心算法每一步选择中都采取在当前状态下最好或最优的选择，通过这种局部最优策略期望产生全局最优解。但是期望是好的，能不能实现是另外一回事了。贪心算法本质上没有从整体最优上加以考虑，并不能保证最终的结果一定是全局最优的。但是相对穷举搜索，搜索效率大大提升。

上个伪代码：

def greedy_decoding(input_ids, max_tokens=300):
    with torch.inference_mode():
        for _ in range(max_tokens):
            outputs = model(input_ids)
            next_token_logits = outputs.logits[:, -1, :]
            next_token = torch.argmax(next_token_logits, dim=-1)
        if next_token == tokenizer.eos_token_id:
            break
        input_ids = torch.cat([input_ids, rearrange(next_token, 'c -> 1 c')], dim=-1)
        generated_text = tokenizer.decode(input_ids[0])
        return generated_text

beam search（束搜索）

简单来说就是每一轮存beam size个候选的句子，这beam size个候选的句子可以得到beam size*vocab size个候选词，在beam size*vocab size个候选词中再选择beam size个继续进入下一轮。

beam search是对greedy search的一个改进算法。相对greedy search扩大了搜索空间，但远远不及穷举搜索指数级的搜索空间，是二者的一个折中方案。

beam search有一个超参数beam size（束宽），设为k。第一个时间步长，选取当前条件概率最大的k个词，当做候选输出序列的第一个词。之后的每个时间步长，基于上个步长的输出序列，挑选出所有组合中条件概率最大的k个，作为该时间步长下的候选输出序列。始终保持k个候选。最后从k个候选中挑出最优的。

还是以上面的任务为例，假设k=2，我们走一遍这个搜索流程。

第一个时间步长：如下图所示，I和H的概率是top2，所以第一个时间步长的输出的候选是I和H，将I和H加入到候选输出序列中。

beam search 第一个时间步长

第2个时间步长：如下图所示，以I开头有三种候选{II, IH, IU}，以H开头有三种候选{HI, HH, HU}。从这6个候选中挑出条件概率最大的2个，即IH和HI，作为候选输出序列。

beam search 第二个时间步长

第3个时间步长：同理，以IH开头有三种候选{IHI, IHH, IHU}，以HI开头有三种候选{HII, HIH, HIU}。从这6个候选中挑出条件概率最大的2个，即IHH和HIU，作为候选输出序列。因为3个步长就结束了，直接从IHH和IHU中挑选出最优值IHU作为最终的输出序列。

beam search 第三个时间步长

beam search不保证全局最优，但是比greedy search搜索空间更大，一般结果比greedy search要好。
greedy search 可以看做是 beam size = 1时的 beam search。

 
import torch
import torch.nn.functional as F
 
def beam_search(LM_prob,beam_size = 3):
    batch,seqlen,vocab_size = LM_prob.shape
    # 对LM_prob取对数，LM_prob: [batch,seqlen,vocab_size]
    # 这个LM_prob不好的地方在于把LM生成的seqlen都放到了LM_prob这一个tensor里了，seqlen的每一个长度代表一次结果，而vocab_size注意不是embedding，代表是vocab里哪个词
    log_LM_prob = LM_prob.log()
    #先选择第0个位置的最大beam_size个token，log_beam_prob与indices的shape为[batch,beam_size]
    log_beam_prob, indices = log_LM_prob[:,0,:].topk(beam_size,sorted = True)
    # topk这个函数刚好是batch内每个元素中所有beam_size之中选topk，log_beam_prob的shape是[batch,beam_size]
    indices = indices.unsqueeze(-1) # indices的shape是[batch,beam_size,1]
    #对每个长度进行beam search
    for i in range(1,seqlen):
        #每个beam的可能产生的概率，由于是log的概率可以直接加
        # log_beam_prob.unsqueeze(-1)的shape是[batch,beam_size,1]
        # log_LM_prob[:,i,:].unsqueeze(1).repeat(1,beam_size,1)的shape是[batch,beam_size,vocab_size]
        log_beam_prob = log_beam_prob.unsqueeze(-1) + log_LM_prob[:,i,:].unsqueeze(1).repeat(1,beam_size,1)
        #选择当前步概率最高的token，此时的log_beam_prob的shape是[batch,beam_size,vocab_size]
        log_beam_prob, index = log_beam_prob.view(batch,-1).topk(beam_size,sorted = True)
        #下面的计算：beam_id选出新beam来源于之前的哪个beam;index代表真实的token id
        #beam_id,index的shape是[batch,beam_size]
        beam_id = index//vocab_size
        index = index%vocab_size
        mid = torch.Tensor([])
        #对batch内每个样本循环，选出beam的同时拼接上新生成的token id
        for j,bid,idx in zip(range(batch),beam_id,index):
            x = torch.cat([indices[j][bid],idx.unsqueeze(-1)], -1) # x.shape是[beam_size, cur_seq_len]
            mid = torch.cat([mid,x.unsqueeze(0)], 0) # mid这时候被拼成[cur_batch, beam_size, cur_seq_len]
        indices = mid #这时候的indices.shape是[batch, beam_size, cur_seq_len]
        print('indices.shape={}'.format(indices.shape))
    return indices, log_beam_prob
 
if __name__=='__main__':
    # 建立一个语言模型 LM_prob (batch,seqlen,vocab_size)，这个纯属偷懒，相当于省去每次decode的工作
    # (batch_i,seqlen_i,vocab_size_i)表示第batch_i个batch中第
    LM_prob = F.softmax(torch.randn([32,20,1000]),dim = -1)
    #最终返回每个候选，以及每个候选的log_prob，shape为(batch,beam_size,seqlen)
    indices,log_prob = beam_search(LM_prob,beam_size = 3)
    # print(indices)

Top-K sampling

简单说beam search的缺点是high quality human language does not follow a distribution of high probability next words，需要一定的随机性，所以有了采样。这就是top-k sampling：在解码的每个时间步从前k个概率最大的词中按它们的概率进行采样。但top-k sampling中k的选择是个难题，选大了可能会采样出长尾词，导致语句不通顺，选小了又退化成了Beam Search。

def top_k_sampling(input_ids, max_tokens=100, top_k=50, temperature=1.0):
    for _ in range(max_tokens):
    with torch.inference_mode():
        outputs = model(input_ids)
        next_token_logits = outputs.logits[:, -1, :]
        top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k)
        top_k_probs = F.softmax(top_k_logits / temperature, dim=-1)
        next_token_index = torch.multinomial(top_k_probs, num_samples=1)
        next_token = top_k_indices.gather(-1, next_token_index)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
        generated_text = tokenizer.decode(input_ids[0])
    return generated_text

Nucleus sampling（top-p sampling）

简单说就是把上面top-k sampling里top k个词换成了top p的概率分布，在每个时间步，头部的几个词的出现概率已经占据了绝大部分概率空间，把这部分核心词叫做nucleus，这个名字起得有点唬人，叫Core Sampling可能更直观些 (但不fancy)

def top_p_sampling(input_ids, max_tokens=100, top_p=0.95):
    with torch.inference_mode():
        for _ in range(max_tokens):
            outputs = model(input_ids)
            next_token_logits = outputs.logits[:, -1, :]
            sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
            sorted_probabilities = F.softmax(sorted_logits, dim=-1) 
            cumulative_probs = torch.cumsum(sorted_probabilities, dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 0] = False 
            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            next_token_logits.scatter_(-1, indices_to_remove[None, :], float('-inf'))
            probs = F.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
        generated_text = tokenizer.decode(input_ids[0])
    return generated_text

Temperature sampling

本质上就是在 Softmax 函数上添加了温度（T）这个参数，以下截图自大模型文本生成——解码策略（Top-k & Top-p & Temperature） - 知乎

import torch
import torch.nn.functional as F

def temperature_sampling(logits, temperature=1.0):
    logits = logits / temperature
    probabilities = F.softmax(logits, dim=-1)
    sampled_token = torch.multinomial(probabilities, 1)
    return sampled_token.item()