从零开始实现大语言模型（十二）：文本生成策略

1. 前言

大语言模型GPTModel通过多轮推理生成连续自然语言文本，每轮推理仅生成一个token。对输入文本做tokenization，将输入文本转换成包含num_tokens个token ID的列表，并输入大语言模型GPTModel，可以得到num_tokens个维度为vocabulary_size的logits向量，第 $i$ 个logits向量是大语言模型根据前 $i$ 个token预测生成的下一个token的概率分数向量，logits向量中的第 $k$ 个概率分数值越大，表明大语言模型预测生成的下一个token的ID为 $k$ 的概率越高。使用softmax函数将最后一个logits向量归一化，使最后一个logits向量每个分量的值均介于0到1之间，所有分量之和等于1，可以得到大语言模型根据输入文本预测生成的下一个token的概率分布。

本文介绍大语言模型GPTModel预测生成连续自然语言文本的流程，以及4种从概率分布中选择下一个token的策略，并实现文本生成函数generate_text。

2. 文本生成流程

大语言模型GPTModel通过多轮推理生成连续自然语言文本，每轮推理仅生成一个token。如下图所示，对输入文本Hello, I am做tokenization，将其转换成包含4个token ID的列表[15496, 11, 314, 716]，并输入大语言模型GPTModel，预测生成ID为257的下一个token a。第2轮推理会将第1轮推理生成的token a添加到输入文本序列，得到包含5个token ID的列表[15496, 11, 314, 716, 257]，并输入大语言模型GPTModel，预测生成ID为2746的下一个token model。依此类推，第6轮推理会将前5轮推理生成的token全部添加到输入文本序列，并将相应token ID列表输入大语言模型GPTModel，最终构造出文本序列Hello, I am a model ready to help.。

3. 文本生成策略

3.1 Greedy Decoding

上述文本生成流程中每轮推理会将包含num_tokens个token ID的列表输入大语言模型GPTModel。根据前文从零开始实现大语言模型（十一）：构建大语言模型GPTModel可知，大语言模型GPTModel会输出num_tokens个维度为vocabulary_size的logits向量，第 $i$ 个logits向量是大语言模型根据前 $i$ 个token预测生成的下一个token的概率分数向量。logits向量中的第 $k$ 个概率分数值越大，表明大语言模型预测生成的下一个token的ID为 $k$ 的概率越高。使用softmax函数将最后一个logits向量归一化，使最后一个logits向量每个分量的值均介于0到1之间，所有分量之和等于1，可以得到大语言模型根据输入文本预测生成的下一个token的概率分布。

Greedy Decoding是一种最简单直接的从概率分布中选择下一个token的策略，其会从大语言模型每轮推理生成的下一个token的概率分布中选择最大概率值对应的index作为预测生成的下一个token的ID。如下图所示，对输入文本Hello, I am做tokenization，将相应token ID列表输入大语言模型GPTModel，并使用softmax函数将大语言模型输出的最后一个logits向量归一化，得到大语言模型根据输入文本Hello, I am预测生成的下一个token的概率分布。Greedy Decoding选择下一个token的概率分布中最大概率值对应的index257作为该轮推理预测生成的下一个token的ID。

可以使用如下代码基于上述大语言模型文本生成策略Greedy Decoding实现大语言模型文本生成函数generate_text_greedy。首先使用tokenizer.encode方法对输入文本做tokenization，将输入文本text转换成包含num_tokens个token ID的列表。在每轮for循环中，使用大语言模型model推理输出num_tokens个维度为vocabulary_size的logits向量，并使用torch.softmax函数将最后一个logits向量归一化，得到下一个token的概率分布。最后使用torch.argmax函数从概率分布中选择最大概率值对应的index作为该轮推理预测生成的下一个token的ID。使用torch.cat方法将token ID列表与预测生成的下一个token的ID拼接起来，构造下一轮推理的输入。执行max_new_tokens轮推理，共生成max_new_tokens个token ID。最后使用tokenizer.decode方法将生成的token ID列表解码，得到大语言模型生成的自然语言文本：

import torchdef generate_text_greedy(model, start_context, max_new_tokens, context_size, tokenizer, stop_ids=None, compact_format=False
):model.eval()idx = tokenizer.encode(start_context, allowed_special=tokenizer.special_tokens_set)idx_tensor = torch.tensor(idx).unsqueeze(0)for _ in range(max_new_tokens):idx_cond = idx_tensor[:, -context_size:]with torch.no_grad():logits = model(idx_cond)logits = logits[:, -1, :]probas = torch.softmax(logits, dim=-1)idx_next = torch.argmax(probas, dim=-1, keepdim=True)if stop_ids is not None and idx_next in stop_ids:breakidx_tensor = torch.cat((idx_tensor, idx_next), dim=1)idx_tensor = idx_tensor.squeeze(0)decoded_text = tokenizer.decode(idx_tensor.tolist())if compact_format:decoded_text = decoded_text.replace("\n", " ")return decoded_text

3.2 Temperature Scaling

对输入文本做tokenization，将相应token ID列表输入大语言模型，并使用softmax函数将大语言模型推理输出的最后一个logits向量归一化，可以得到下一个token的概率分布。将同一自然语言文本多次输入大语言模型，每次推理预测生成的logits向量会完全相同，上述文本生成策略Greedy Decoding使用torch.argmax函数从概率分布中选择最大概率值对应的index作为当次推理预测生成的下一个token的ID，因此多轮推理会预测生成完全相同的自然语言文本序列。

可以使用torch.multinomial函数根据下一个token的概率分布做随机抽样，使每轮推理预测生成的下一个token的多样性更强，大语言模型多轮推理可以预测生成多种多样的自然语言文本序列。

如下面的代码所示，构造一个仅包含9个token的词汇表vocabulary，并假设大语言模型根据输入文本every effort moves you推理输出的最后一个logits向量为[4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]。使用torch.softmax函数将该logits向量归一化，得到下一个token的概率分布probas，并使用torch.multinomial函数根据概率分布做随机抽样，打印输出当次推理预测生成的下一个token：

torch.manual_seed(123)vocabulary = {"closer": 0, "every": 1, "effort": 2, "forward": 3, "inches": 4, "moves": 5,  "pizza": 6, "toward": 7, "you": 8
}
inverse_vocabulary = list(vocabulary.keys())next_token_logits = torch.tensor([4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79])
probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.multinomial(probas, num_samples=1).item()
print(inverse_vocabulary[next_token_id])

执行上面代码，打印结果如下：

forward

可以使用如下代码循环执行1000次上述使用torch.multinomial函数根据概率分布做随机抽样的流程，打印输出所有生成的下一个token及其次数：

def print_sampled_tokens(probas):torch.manual_seed(123)sample = [torch.multinomial(probas, num_samples=1).item() for i in range(1000)]sampled_ids = torch.bincount(torch.tensor(sample))for i, freq in enumerate(sampled_ids):print(f"{freq} x {inverse_vocabulary[i]}")print_sampled_tokens(probas)

执行上面代码，打印结果如下：

73 x closer
0 x every
0 x effort
582 x forward
2 x inches
0 x moves
0 x pizza
343 x toward

从上面的打印结果可知，虽然logits向量中最大概率分数值对应的token forward被随机抽样选中的次数最多（582次），但是也可能会生成其他token。使用torch.multinomial函数根据下一个token的概率分布做随机抽样，可以让大语言模型推理生成自然语言文本序列的多样性更强。

Temperature Scaling是一种将大语言模型推理输出的logits向量除以一个大于0的温度常数temperature，以调整softmax函数输出下一个token的概率分布的方法。当温度常数temperature等于1，会在使用softmax函数将最后一个logits向量归一化之前，先将该logits向量除以1，即如果温度常数temperature等于1，softmax函数输出下一个token的概率分布会与不使用Temperature Scaling的情况完全相同。温度常数temperature越大，随机性会越强，下一个token的概率分布中各个token的概率值会越接近。温度常数temperature越小，随机性会越弱，下一个token的概率分布中概率值会更集中于部分概率分数值更大的token。

使用Temperature Scaling方法对logits向量做变换，将logits向量next_token_logits分别除以三个temperature常数[1, 0.1, 5]，并使用torch.softmax函数分别将三个经过变换的logits向量归一化，打印输出使用不同temperature时softmax函数输出的下一个token的概率分布：

import matplotlib.pyplot as pltdef softmax_with_temperature(logits, temperature):scaled_logits = logits / temperaturereturn torch.softmax(scaled_logits, dim=0)temperatures = [1, 0.1, 5]
scaled_probas = [softmax_with_temperature(next_token_logits, t) for t in temperatures]
x = torch.arange(len(vocabulary))
bar_width = 0.15
fig, ax = plt.subplots(figsize=(5, 3))
for i, t in enumerate(temperatures):rects = ax.bar(x + i * bar_width, scaled_probas[i], bar_width, label=f"Temperature = {t}")
ax.set_ylabel("Probability")
ax.set_xticks(x)
ax.set_xticklabels(vocabulary.keys(), rotation=90)
ax.legend()
plt.tight_layout()
plt.show()

执行上面代码，生成使用不同temperature时softmax函数输出下一个token的概率分布图像如下：

根据上面的图像可知，如果不使用Temperature Scaling（等同于温度常数temperature等于1），使用torch.multinomial函数根据下一个token的概率分布做随机抽样，生成的下一个token大约有60%的概率是forward。温度常数temperature越小，生成的下一个token为logits向量中概率分数值最大的token forward的概率越高。温度常数temperature越大，softmax函数输出的下一个token的概率分布中各个token的概率值会越接近，大语言模型生成的自然语言文本序列的多样性会更强。

3.3 Top-k Sampling

使用torch.multinomial函数从经过Temperature Scaling策略变换的下一个token的概率分布中做随机抽样，可以让大语言模型多轮推理预测生成多种多样的自然语言文本序列。温度常数temperature越大，每轮推理softmax函数输出的下一个token的概率分布中各个token的概率值会越接近，大语言模型生成的自然语言文本序列的多样性会更强。

使用Temperature Scaling策略可以显著增强大语言模型的创造性，但是也可能会让大语言模型生成有语法错误或者完全无意义的内容。从3.2中绘制的当温度常数temperature等于5时softmax函数输出下一个token的概率分布图像可知，大约有4%的概率会生成every effort moves you pizza这样没有意义的自然语言文本。

Top-k Sampling是一种限制大语言模型每轮推理输出token的候选范围，以提升大语言模型生成的自然语言文本序列质量的方法。如下图所示，Top-k Sampling会保留logits向量中概率分数值较大的top_k个token作为当次推理输出token的候选集，并使用mask屏蔽所有不在候选集中token的概率分数，使所有不在候选集中token的概率分数值全都为 $-\infty$ ，从而让softmax函数输出的下一个token的概率分布中所有候选集中的token的概率值之和等于1，不在候选集中的所有token的概率值全都为0。

如下面的代码所示，可以使用torch.topk函数从logits向量next_token_logits获取top_k个概率分数值及其相应token的ID，并使用torch.where函数将logits向量中其他token对应的概率分数值全部变成 $-\infty$ ：

top_k = 3top_logits, top_pos = torch.topk(next_token_logits, top_k)
new_logits = torch.where(condition=next_token_logits < top_logits[-1],input=torch.tensor(float("-inf")),other=next_token_logits
)print("Top logits:", top_logits)
print("Top positions:", top_pos)
print("New logits:", new_logits)

执行上面代码，打印结果如下：

Top logits: tensor([6.7500, 6.2800, 4.5100])
Top positions: tensor([3, 7, 0])
New logits: tensor([4.5100,   -inf,   -inf, 6.7500,   -inf,   -inf,   -inf, 6.2800,   -inf])

使用torch.softmax函数将经过Top-k Sampling策略变换的logits向量归一化，打印输出下一个token的概率分布：

topk_probas = torch.softmax(new_logits, dim=0)
print(topk_probas)

执行上面代码，打印结果如下：

tensor([0.0615, 0.0000, 0.0000, 0.5775, 0.0000, 0.0000, 0.0000, 0.3610, 0.0000])

3.4 Top-p Sampling

Top-p Sampling是另一种限制大语言模型每轮推理输出token的候选范围，以提升大语言模型生成的自然语言文本序列质量的方法。使用softmax函数将logits向量归一化，可以得到下一个token的概率分布。Top-p Sampling会保留概率分布中概率值之和刚好不小于top_p的 $n$ 个较大概率值对应的token作为当次推理输出token的候选集，并使用3.3中所述的mask方法屏蔽所有不在候选集中token的概率分数，使softmax函数输出的下一个token的概率分布中所有候选集中的token的概率值之和等于1，不在候选集中的所有token的概率值全都为0。

如下面的代码所示，可以使用torch.sort函数将torch.softmax函数输出的下一个token的概率分布中的概率值按从大到小排序，并使用torch.cumsum函数计算累积概率cumulative_probas。torch.searchsorted函数可以返回cumulative_probas中累积概率刚好不小于top_p的索引位置cutoff_index，进而通过next_token_logits[sorted_indices[cutoff_index]]得到累积概率刚好不小于top_p的 $n$ 个较大概率分数集中最小的概率分数值。最后可以使用torch.where函数将logits向量中不在候选集中的所有token的概率分数值全都变成 $-\infty$ ：

top_p = 0.9probas = torch.softmax(next_token_logits, dim=0)
sorted_probas, sorted_indices = torch.sort(probas, descending=True)
cumulative_probas = torch.cumsum(sorted_probas, dim=0)
cutoff_index = torch.searchsorted(cumulative_probas, top_p)new_logits = torch.where(condition=next_token_logits < next_token_logits[sorted_indices[cutoff_index]],input=torch.tensor(float("-inf")),other=next_token_logits
)print("New logits:", new_logits)

执行上面代码，打印结果如下：

New logits: tensor([  -inf,   -inf,   -inf, 6.7500,   -inf,   -inf,   -inf, 6.2800,   -inf])

使用torch.softmax函数将经过Top-p Sampling策略变换的logits向量归一化，打印输出下一个token的概率分布：

topp_probas = torch.softmax(new_logits, dim=0)
print(topp_probas)

执行上面代码，打印结果如下：

tensor([0.0000, 0.0000, 0.0000, 0.6154, 0.0000, 0.0000, 0.0000, 0.3846, 0.0000])

4. 文本生成函数

可以结合上述4种文本生成策略实现大语言模型文本生成函数generate_text。修改3.1中实现的基于Greedy Decoding策略的大语言模型文本生成函数generate_text_greedy，在每轮for循环中，使用3.3中所述Top-k Sampling策略及3.4中所述Top-p Sampling策略共同确定当次推理预测输出token的候选集，并使用torch.where函数对最后一个logits向量做变换，使所有不在候选集中token的概率分数值全都为 $-\infty$ 。使用3.2中所述Temperature Scaling策略，将经过变换的最后一个logits向量除以大于0的温度常数temperature，并使用torch.softmax函数将经过上述策略变换的logits向量归一化，得到下一个token的概率分布。最后使用torch.multinomial函数根据概率分布做随机抽样，得到当次推理预测生成的下一个token的ID。具体代码如下所示：

def generate_text(model, start_context, max_new_tokens, context_size, tokenizer,temperature=0.0, top_k=None, top_p=None, stop_ids=None, compact_format=False
):model.eval()idx = tokenizer.encode(start_context, allowed_special=tokenizer.special_tokens_set)idx_tensor = torch.tensor(idx).unsqueeze(0)for _ in range(max_new_tokens):idx_cond = idx_tensor[:, -context_size:]with torch.no_grad():logits = model(idx_cond)logits = logits[:, -1, :]min_val = Noneif top_k is not None:top_logits, _ = torch.topk(logits, top_k, dim=-1)min_val = top_logits[:, -1]if top_p is not None:probas = torch.softmax(logits.squeeze(0), dim=-1)sorted_probas, sorted_indices = torch.sort(probas, descending=True, dim=-1)cumulative_probas = torch.cumsum(sorted_probas, dim=-1)cutoff_index = torch.searchsorted(cumulative_probas, top_p)min_val_topp = logits[:, sorted_indices[cutoff_index]]min_val = min_val_topp if min_val is None else torch.max(min_val, min_val_topp)if min_val is not None:logits = torch.where(condition=logits < min_val,input=torch.tensor(float("-inf")).to(logits.device),other=logits)if temperature > 0.0:logits = logits / temperatureprobas = torch.softmax(logits, dim=-1)idx_next = torch.multinomial(probas, num_samples=1)else:idx_next = torch.argmax(logits, dim=-1, keepdim=True)if stop_ids is not None and idx_next in stop_ids:breakidx_tensor = torch.cat((idx_tensor, idx_next), dim=1)idx_tensor = idx_tensor.squeeze(0)decoded_text = tokenizer.decode(idx_tensor.tolist())if compact_format:decoded_text = decoded_text.replace("\n", " ")return decoded_text

5. 结束语

对输入文本做tokenization，将输入文本转换成包含num_tokens个token ID的列表，并输入大语言模型GPTModel，可以得到num_tokens个维度为vocabulary_size的logits向量。使用上述文本生成函数generate_text可以利用大语言模型GPTModel进行多轮推理生成连续自然语言文本。至此，使大语言模型GPTModel具备GPT-3那样的文本补全及小样本学习能力，只剩下最后一朵乌云：预训练大语言模型GPTModel！