深入理解束搜索（Beam Search）

本文主要内容参考以下博客：

• 【AI知识点】束搜索（Beam Search）
• 什么是 Beam Search？深入理解模型生成策略

束搜索（Beam Search） 是一种 常用于序列生成任务的启发式搜索算法，广泛应用于自然语言处理中的机器翻译、文本生成等任务。它是一种 平衡了搜索效率和搜索质量的方法，相比于贪心搜索，它能找到更优的解，但计算复杂度比穷举搜索更低。

1. 核心思想

束搜索是一种宽度受限的最优搜索方法，在生成序列时，它会同时保留多个最优候选，而不是只关注当前最优的选择。具体来说，束搜索会保留固定数量的候选序列，这被称为束宽（beam width or beam size），并在每一步扩展这些候选序列，最终从中选择得分最高的一个作为结果。

2. 工作原理

1. 初始化：从序列生成的初始状态（通常是句子开头的一个标记，如 $<start>$）开始，模型为每个候选生成概率。
2. 束宽控制：根据概率值，选出前 $k$ 个最优的候选，$k$ 就是束宽，这些候选将被进一步扩展。
3. 逐步生成：对每个选中的候选项，模型继续生成下一个词的概率分布，并扩展所有候选。再从所有扩展候选中选出 $k$ 个最优序列。
4. 终止条件：直到所有序列生成完毕（例如，遇到 $<end>$ 标记），最终选择得分最高的序列作为生成的结果。

3. 处理过程

3.1 具体示例

假设我们有一个简单的语言模型要生成句子，模型根据之前的词生成下一个词的概率。以下是一个生成句子的例子，束宽为 2（即每一步保留两个最优候选）：

初始状态：
模型预测三个可能的词：

“我” (0.5)
“你” (0.3)
“他” (0.2)

选择得分最高的前两个候选，“我”和“你”。

第一步扩展：

对于“我”，模型预测：

“喜欢” (0.6)
“讨厌” (0.4)

对于“你”，模型预测：

“喜欢” (0.7)
“讨厌” (0.3)

选择得分最高的前两个候选：

“我喜欢” (0.5 * 0.6 = 0.3)
“你喜欢” (0.3 * 0.7 = 0.21)

第二步扩展：
对于“我喜欢”和“你喜欢”，模型继续生成下一个词。最终生成：

“我喜欢吃” (得分更高)
“你喜欢跑”

3.2 可视化

Beam Search Visualizer：非常棒的Beam Search交互式项目

可视化如下：

4. 束宽的影响

• 较小束宽（如 $k=1$）：类似于贪心搜索，只保留当前最优的选择，可能导致找到的解局部最优但全局较差。
• 较大束宽：增加搜索范围，保留更多可能的候选序列，找到的解更接近全局最优，但计算成本增加。

5. 应用场景

束搜索在以下场景中应用广泛：

• 机器翻译：生成目标语言句子时，使用束搜索在每一步保留多个翻译候选，选择最优翻译。
• 文本生成：如自动写作、对话生成等，束搜索可以帮助生成更连贯和合理的句子。
• 语音识别：在识别语音时，束搜索用于从多个可能的词序列中找到最可能的识别结果。

6. 计算复杂度

束搜索的计算复杂度主要由以下几个因素决定：

• 束宽 $k$：即每一步保留的候选序列数量。
• 序列长度 $T$：即需要生成的目标序列的长度。
• 每一步生成的候选项数 $V$：通常对应于模型的词汇表大小（Vocabulary size）。

在每一步，束搜索从 $k$ 个候选序列中分别扩展 $V￥ 个新词，因此单步的计算复杂度为 $O(k\cdot V)$。在整个生成序列的过程中，束搜索需要进行 $T$ 步扩展，因此，束搜索的总计算复杂度为：$O(k\cdot V\cdot T)$

假设我们有以下参数：

• 词汇表大小 V=10,000
• 序列长度 T=20
• 束宽 k=5

那么，束搜索的总计算复杂度将是：$O(5\cdot 10,000\cdot 20)=O(1,000,000)$

总结

束搜索是一种在序列生成任务中用于平衡搜索效率和准确性的算法，它通过同时保留多个候选序列，能够避免局部最优解问题。其束宽决定了搜索范围和计算复杂度，广泛应用于机器翻译、文本生成和语音识别等任务。

参考资料

• 【AI知识点】束搜索（Beam Search）
• 什么是 Beam Search？深入理解模型生成策略