生成式语言模型常见的采样方式有贪心采样、top k采样、top p采样和beam search采样。 top k采样和top p采样都是对候选token作“截断”，仅考虑概率最大的若干个token的方法。top k只从概率最高的k个候选token中进行随机选择。top p同样优先考虑高概率的token，直到候选token的概率累积达到p。贪心采样可以视为top k和top p的特例，即总是只考虑概率最高的那个token。
无论是贪心、top k还是top p，它们都只探索了一条生成序列，容易落入局部最优。Beam Search的特点是通过在搜索过程中同时考虑概率最高的若干条可能序列，这样就扩大了搜索范围。
以我个人的实际项目经验来看，只要用上beam search，一般模型的评测得分都能小涨几个点。在一些对输出字符串有限制的任务中，我们还可以在beam search内引入约束，限制模型的生成。Beam search是一个基础、方便、能够即插即用的改善生成质量的生成策略。这篇文章就来简单记录下beam search的原理。
假设下面的model存储了每个token被生成的概率，我们讨论如何基于这个模型实现beam search采样。
显示初始化代码
import math
import numpy as np
import numpy as np
from dataclasses import dataclass
from typing import List
from numbers import Number
import functools
import random
random.seed(0)
model = {
'a': 0.7,
'<eos>': 0.3,
}这是一个简化了的语言模型——在实际应用中，语言模型的概率分布会随上下文变化，但这里我们假设token产生的概率是与历史token无关的。
Beam search的要点是，在每次生成下一个token时列举出所有可能，然后只保留使整体序列概率最大的那n个。
为了理解beam search，我们先手推一遍基于model模型的beam search生成过程，然后再编写程序与手推的结果进行比对。
beam search的两个基本参数是beam_width和tokens_to_gen. beam_width表示在生成时要考虑多少个候选序列，而tokens_to_gen表示要生成多少个字符。
假设beam_width = 2，我们在空字符串的基础上生成3个token。手推生成beam search的生成过程，我们会得到这样的一个生成树：
flowchart LR
subgraph "第0步"
A
end
subgraph "第1步"
A["空串, 1"] --> B("a, 1 * 0.7 = 0.7")
A --> C("<eos>, 1 * 0.3 = 0.3")
end
subgraph "第2步"
B --> D("aa, 0.7 * 0.7 = 0.49")
B --> E("a<eos>, 0.3 * 0.7 = 0.21")
C --> F("<eos>, 0.3")
end
subgraph "第3步"
D --> G("aaa, 0.49 * 0.7 = 0.343")
D --> H("aa<eos>, 0.49 * 0.3 = 0.147")
F --> I("<eos>, 0.3")
end
上图以树的形式展示了beam search的生成过程。每个节点都以“字符串，概率”的形式表示了当前生成的文本和对应的概率。每一步概率最大的两个节点会被选中，在下一步生成后续token。在计算过程中需要注意，一旦<eos>token被生成，那么它不会再产生后续token，其概率也不会再被改变。每生成一个token，当前生成结果的概率便等于上一个token序列的概率乘以生成token的概率。
在树状图的最后一层，我们可以看到算法的最终生成结果。取出概率最大的两个生成结果，分别得到aaa和<eos>. 下面就让我们编写对应的beam search程序，并对照检查结果是否正确。
1 Beam Search代码实现
我们先构造一个BeamSearchCandidate类用于表示beam search的生成结果。这个类储存了当前生成的token和对应的概率。
@functools.total_ordering
@dataclass
class BeamSearchCandidate:
tokens: List[str]
logprob: Number
def __repr__(self):
return f'{"".join(self.tokens)}:{math.exp(self.logprob):.2f}'
def __ge__(self, other: 'BeamSearchCandidate'):
return self.logprob > other.logprob实践中，我们经常存储概率的对数，即代码中的logprob，这样能将概率的相乘转化为对数概率的相加，减少计算代价，改善数值稳定性问题。
在理解算法运行过程的基础上，程序的实现并不困难。同样，以model这个简单模型为例，beam search算法实现如下：
def simple_top_k(arr, k):
'''一种简单的top k算法。'''
arr.sort()
return arr[-k:]
import tqdm
def beam_search(
model,
input_sequence:str,
beam_width:int,
tokens_to_gen:int,
top_k=simple_top_k, # ←后续我们会讨论top k算法的优化
) -> List[BeamSearchCandidate]:
candidates = [BeamSearchCandidate([input_sequence], 0.)]
for i in tqdm.tqdm(range(tokens_to_gen)):
next_candidates = []
for cdd in candidates:
if len(cdd.tokens) and cdd.tokens[-1] == '<eos>':
next_candidates.append(cdd)
else:
for c, lp in zip(model.tokens, model.logprob(cdd.tokens)):
next_candidates.append(
BeamSearchCandidate(
cdd.tokens + [c], # 计算每个候选语句的概率
cdd.logprob + lp
)
)
# 每一步都只保留概率最高的那些候选语句
candidates = top_k(next_candidates, k=beam_width)
return candidates从beam search的运算结果可以看到，其生成结果与我们手动计算的一致。
ret = beam_search(model, '', beam_width=2, tokens_to_gen=3)
print(ret)100%|██████████| 3/3 [00:00<00:00, 2092.62it/s][<eos>:0.30, aaa:0.34]接下来本文记录笔者实践中遇到的一些问题和解决思路。
2 Top K算法优化
Beam search算法中，选择概率最大的若干个候选token序列是一个比较耗时的操作。以上beam search函数使用了排序算法来找到前\(k\)个概率最大的候选文本，这一操作的时间复杂度为\(O(n\log n)\)，\(n\)为总的候选token数量。因为大模型的词表比较大，所以改进beam search所用的top k算法能带来不少加速。
第一个思路是对选择排序进行改造。在排序的过程中，使其在完成k个数字的排序时停止。这样就获得了一种简单的top k选择算法，此类方法的时间复杂度是\(O(nk)\)，\(k\)为beam width。
另一种思路是基于优先队列（堆）实现top k算法。我们可以维护一个大小为\(k\)的堆，在遍历过所有\(n\)个候选token后，堆中留下的便是概率最大的\(k\)个预测序列。维护堆的复杂度为\(O(\log k)\)，共维护\(n\)次，因此总的时间复杂度为\(O(n\log k)\).
最后，一种非常高效的，但知道的人比较少的top k算法是快速选择算法（quick select），其时间复杂度为\(O(n)\). 显然这是top k算法所能达到的最优时间复杂度——因为不管哪种算法也要把原始输入序列逐个过一遍。
与前面几种top k算法相比，快速选择算法的缺点在于其无法保证返回的k个结果的有序性。即quick select方法以牺牲结果的有序性为代价改进了其时间复杂度。
3 数值精度优化
- 大模型的词表一般很大，这就导致所有token的生成概率可能会整体偏低；
- 模型内部实现上并不是直接返回\([0, 1)\)之间的概率数值，而是返回范围在\((-\infty, \infty)\)内的logits；
- 为了加速推理和训练，人们经常使用
float16等低精度浮点型。
以上因素的共同作用下，beam search很容易遇到数值溢出的问题。在实践中，可以利用将logits同时加上或者减去任意数，不影响概率大小的特性来调控logits数值的范围，尽量避免溢出。这一原理我在之前的《Softmax原理》一文中也有讲解。
4 长度归一化
模型有时候可能会倾向于产生较长或较短的文本，这时候对概率作长度归一化是一种可以考虑的手段。
在未引入长度归一化时，beam search选取候选token的条件为 \[ \arg\max \sum_{t=1}^{T_y}\log P(y_t|X, y_1, y_2, \dots, y_{t-1}), \] 其中\(T_y\)为序列长度。
在引入长度归一化后，条件变为： \[ \arg\max \frac{1}{T_y^\alpha} \sum_{t=1}^{T_y}\log P(y_t|X, y_1, y_2, \dots, y_{t-1}). \] 新的条件引入了一个额外的参数\(\alpha\in[0, 1]\)。当\(\alpha\)为0时，beam search的行为相当于不作任何归一化；\(\alpha=1\)则对应著“标准的”归一化操作。如果想要取得一个折中的效果，可以设\(\alpha\)为\((0, 1)\)内的一个数字。
