怎么理解BeamSearch？

在大模型推理中，常会用到BeamSearch，本文就BeamSearch原理与应用理解展开讲解。

一、BeamSearch原理

Beam Search 是一种启发式搜索算法，常用于自然语言处理（NLP）和其他需要生成序列的任务中，比如机器翻译、自动摘要和语音识别，大模型推理等。它是一种改进的贪心算法，旨在平衡计算效率与搜索质量。

基本思想
Beam Search 通过维护一个有限大小的候选集（称为“beam”），在每一步只选择最有希望的若干个候选项，从而避免了暴力搜索空间的指数级增长。简言之，Beam Search 会在每个时间步上选择最佳的“beam width”（宽度）个候选，而不是总是选择最好的单个候选。

二、工作流程

假设要生成一个序列，算法的步骤如下：

1. 初始化：开始时，Beam Search 选择一个初始状态（比如一个空的序列）并评估其可能性。
2. 扩展候选项：每一步，算法生成所有可能的下一个词或状态，并评估这些候选项的得分（通常通过某种概率模型来评估，如语言模型或神经网络的输出）。
3. 选择最佳候选：然后，只保留得分最高的“beam width”个候选项，并丢弃其他候选项。这些候选项将继续向下扩展。
4. 停止条件：直到满足某种停止条件（如生成完毕或达到最大长度）时，算法会停止，最终选择得分最高的候选序列作为结果。

三、举例说明

假设我们正在进行一个简单的语言模型任务，要生成句子“I am a student”。假设我们的候选词在每个时间步都来自一个有限的词汇表。

1. Step 1 - 选择初始词
   我们从“I”开始生成序列。
2. Step 2 - 选择下一个单词
   我们现在需要从“I”出发，选择下一个单词。假设当前候选词和它们的概率（或得分）如下：
   “am”（概率 = 0.6）
   “is”（概率 = 0.2）
   “was”（概率 = 0.1）
   “I”（概率 = 0.1）
   假设我们选择 Beam Width = 2，因此我们保留得分最高的 2 个候选单词：
   选择“am”（得分 0.6）
   选择“is”（得分 0.2）
3. Step 3 - 继续生成下一个单词
   现在，我们有两个候选序列：
   “I am”
   “I is”
   对于每个序列，我们根据模型的得分选择下一个词。假设得分如下：
   对于 “I am”（从“am”继续）：
   “a”（概率 = 0.5）
   “the”（概率 = 0.3）
   “an”（概率 = 0.2）
   对于 “I is”（从“is”继续）：
   “am”（概率 = 0.4）
   “was”（概率 = 0.3）
   “a”（概率 = 0.2）
   我们再次选择 Beam Width = 2，所以保留概率最大的 2 个候选项：
   “I am a”（得分 0.6 × 0.5 = 0.3）
   “I am the”（得分 0.6 × 0.3 = 0.18）
   “I is am”（得分 0.2 × 0.4 = 0.08）
   “I is was”（得分 0.2 × 0.3 = 0.06）
   最终，我们选择前两个候选：
   “I am a”
   “I am the”
4. Step 4 - 决定最后的单词
   接下来，我们继续为每个候选序列添加下一个单词：
   对于 “I am a”（从“a”继续）：
   “student”（概率 = 0.7）
   “teacher”（概率 = 0.2）
   “doctor”（概率 = 0.1）
   对于 “I am the”（从“the”继续）：
   “student”（概率 = 0.6）
   “man”（概率 = 0.3）
   “teacher”（概率 = 0.1）
   继续选择 Beam Width = 2，保留得分最高的两个候选：
   “I am a student”（得分 0.3 × 0.7 = 0.21）
   “I am a teacher”（得分 0.3 × 0.2 = 0.06）
   “I am the student”（得分 0.18 × 0.6 = 0.108）
   “I am the man”（得分 0.18 × 0.3 = 0.054）
   最终，我们选择最好的两个序列：
   “I am a student”
   “I am the student”
   Step 5 - 完成生成

假设我们已经达到了序列的结束条件（比如长度或特殊的结束符），最后选择 “I am a student” 作为最优序列。

如果当 Beam Width = 2，但是有三个概率相等的候选时，我们通常会遇到一个需要做出决策的情况。Beam Search 是基于保留每步得分最好的候选项来进行的，但是如果有多个候选项的概率相同，可以选择随机从这些候选中选出若干个，或者可以选择 按顺序 排列所有候选项，并选出得分相同的前 N 个候选项。

停止生成的条件：

1. 生成的序列有一个预设的最大长度，超过这个长度就停止生成。
2. 在生成句子的任务中，如果当前候选序列已经包含 ，而该序列的后续扩展也不会再生成其他单词，那么就可以提前停止对该序列的扩展，直接输出该序列

四、与贪心算法的对比

• 贪心算法：每一步只选择当前最优的选项，不回溯或考虑全局优化。
• Beam Search：每一步选择多个候选项（由宽度控制），并且可以保持多个候选状态，在后续步骤中进行进一步扩展，从而能够考虑更多的可能性。

如果我们使用 贪心算法，每一步只选择最有可能的下一个词。假设我们在第一步选择 “am” 后，接下来我们会选择“a”作为下一个单词，然后选择“student”。但是这种方法可能会导致不理想的结果，而 Beam Search 则能够考虑多个可能性，最终找到最佳序列。

五、Beam Search 的优缺点

• 优点：
  比贪心算法更好：能够考虑多个候选项，避免早期陷入局部最优解。
  计算效率较高：相对于暴力搜索，Beam Search 大大减少了搜索空间。
• 缺点：
  计算开销较大：随着 Beam Width 增大，计算的复杂度也会增加。
  不能保证全局最优解：由于只保留固定数量的候选项，可能会错过一些较优的解。