Greedy Decoding

Greedy Decoding在每个时间步选择当前条件概率最高的词语作为输出，直到生成结束。在贪婪解码中，生成模型根据输入序列，逐个时间步地预测输出序列中的每个词语。在每个时间步，模型根据当前的隐藏状态和已生成的部分序列计算每个词语的条件概率分布，模型选择具有最高条件概率的词语作为当前时间步的输出。这个词语成为下一个时间步的输入，生成过程持续直到满足某种终止条件，比如生成了指定长度的序列或者生成了特殊的结束标记。

这种方法简单高效，每个时间步只需计算当前条件概率最高的词语，因此计算速度较快。但是由于每个时间步只考虑当前条件概率最高的词语，贪婪解码可能会陷入局部最优解，而无法获得全局最优解。这可能导致生成的文本缺乏多样性或不准确。

尽管贪婪解码存在一些局限性，但它仍然是许多序列生成任务中常用的一种方法，特别是在对速度要求较高或者任务较为简单的情况下。

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, torch.unsqueeze(next_token, 0)], dim=-1)
        generated_text = tokenizer.decode(input_ids[0])
    return generated_text

Beam Search

束搜索（Beam Search）是贪婪解码的一种扩展，通过在每个时间步保留多个候选序列来克服贪婪解码的局部最优问题。

在每个时间步保留概率最高的前几个候选词语，然后在下一个时间步基于这些候选词语继续扩展，直到生成结束。束搜索通过考虑多个候选词语路径，可以在一定程度上增加生成文本的多样性。

在束搜索中，模型在每个时间步会生成多个候选序列，而不是仅选择一个最优序列。模型会根据当前已生成的部分序列和隐藏状态，预测下一个时间步可能的词语，并计算每个词语的条件概率分布。

上图的每一步中，只保留两条最可能的路径(根据beam =2)，而所有其他都被丢弃。此过程将继续进行，直到满足停止条件，该停止条件可以是生成序列结束令牌或达到最大序列长度的模型。最终输出将是最后一组路径中具有最高总体概率的序列。

from einops import rearrange
import torch.nn.functional as F

def beam_search(input_ids, max_tokens=100, beam_size=2):
    device = input_ids.device
    beam_scores = torch.zeros(beam_size, device=device)
    beam_sequences = input_ids.clone()
    active_beams = torch.ones(beam_size, dtype=torch.bool, device=device)
    for step in range(max_tokens):
        outputs = model(beam_sequences)
        logits = outputs.logits[:, -1, :]
        probs = F.softmax(logits, dim=-1)
        top_scores, top_indices = torch.topk(probs.flatten(), k=beam_size, sorted=False)
        beam_indices = top_indices // probs.shape[-1]
        token_indices = top_indices % probs.shape[-1]
        new_sequences = torch.cat([
            beam_sequences[beam_indices],
            token_indices.unsqueeze(-1)
        ], dim=-1)
        beam_sequences = new_sequences
        beam_scores = top_scores
        active_beams = ~(token_indices == tokenizer.eos_token_id)
        if not active_beams.any():
            print("no active beams")
            break
    best_beam = beam_scores.argmax()
    best_sequence = beam_sequences[best_beam]
    generated_text = tokenizer.decode(best_sequence)
    return generated_text

Temperature Sampling

温度参数采样（Temperature Sampling）常用于基于概率的生成模型，如语言模型。它通过引入一个称为“温度”（Temperature）的参数来调整模型输出的概率分布，从而控制生成文本的多样性。

在温度参数采样中，模型在每个时间步生成词语时，会计算出词语的条件概率分布。然后模型将这个条件概率分布中的每个词语的概率值除以温度参数，对结果进行归一化处理，获得新的归一化概率分布。较高的温度值会使概率分布更平滑，从而增加生成文本的多样性。低概率的词语也有较高的可能性被选择；而较低的温度值则会使概率分布更集中，更倾向于选择高概率的词语，因此生成的文本更加确定性。最后模型根据这个新的归一化概率分布进行随机采样，选择生成的词语。

import torch
import torch.nn.functional as F

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

Top-K Sampling

Top-K 采样（在每个时间步选择条件概率排名前 K 的词语，然后在这 K 个词语中进行随机采样。这种方法既能保持一定的生成质量，又能增加文本的多样性，并且可以通过限制候选词语的数量来控制生成文本的多样性。