Beam Search 束搜索

Beam Search（束搜索）是序列生成中最经典的解码算法之一，介于贪心解码的”只顾眼前”与穷举搜索的”算力爆炸”之间，用固定宽度的候选集兼顾输出质量与计算成本。前置阅读：数值优化与数学基础，相关章节见解码策略与采样。

核心直觉

生成文本 = 在一棵不断分叉的树形迷宫里找一条最优路径。

贪心解码（Greedy）= 走迷宫时每到一个岔路口，只看眼前哪条路最近就钻进去，绝不回头。简单快，但一旦走错一步，后面全盘皆错。
穷举搜索（Exhaustive Search）= 把迷宫的每一条分支都走一遍，最后挑最短的那条。绝对最优，但岔路呈指数增长，词表稍大就彻底算不动。
Beam Search（束搜索）= 走迷宫时口袋里始终揣着 k 条”最有希望的探路小分队”，每个岔路口各分队都向四面八方试探一步，然后从所有新路线里留下总分最高的 k 条，淘汰其余。既不像贪心那样孤注一掷，也不像穷举那样铺张浪费。

一句话概括：贪心是 width=1 的束搜索；穷举是 width=无穷大的束搜索；束搜索就是介于两者之间的折中。

原理详解

从贪心说起

自回归模型（Autoregressive Model，即每步基于已生成内容预测下一个 token 的模型，如 GPT、机器翻译的 Decoder）生成文本时，每一步都在整个词表上输出一个概率分布，即给定已生成的前缀，预测第 t 步 token 的概率。贪心解码的做法是：每步取概率最大的那个 token（argmax），把它拼回输入继续预测下一个。整个序列的联合概率就是各步条件概率的乘积。

用数学语言描述：给定已生成前缀 y_{1:t-1}，模型输出词表中每个 token w_i 的条件概率 p(w_i | y_{1:t-1})。贪心解码每步选择：

y_t = \arg\max_{w \in V} p(w \mid y_{1:t-1})

整个序列的联合概率为各步条件概率的乘积：

P(y_1, y_2, \ldots, y_T) = p(y_1) \times p(y_2 \mid y_1) \times \cdots \times p(y_T \mid y_1, \ldots, y_{T-1}) = \prod_{t=1}^{T} p(y_t \mid y_{1:t-1})

贪心解码的致命问题是局部最优陷阱：第一步挑了概率最高的词，可能把后面引入了一个更差的分支，而那条”起步概率略低但后续更顺”的路径被永远丢弃了。

举一个具体例子说明贪心的盲区：

词表（简化）: { "我", "喜欢", "是", "的", "爱", "程序" }

步骤 1:
  p("我" | <BOS>) = 0.50   ← 贪心选这个
  p("爱" | <BOS>) = 0.45

步骤 2（假设贪心选了"我"）:
  p("是" | "我") = 0.30    → 联合概率 = 0.50 × 0.30 = 0.150
  p("喜欢" | "我") = 0.25  → 联合概率 = 0.50 × 0.25 = 0.125

步骤 2（假设选了"爱"）:
  p("程序" | "爱") = 0.60  → 联合概率 = 0.45 × 0.60 = 0.270  ← 更高！

贪心在第 1 步选了”我”（0.50 > 0.45），但”爱程序”这条路径的联合概率（0.270）远高于贪心路径（0.150）。Beam Search 会在第 1 步同时保留”我”和”爱”两条候选，从而在后续步骤中发现更优的完整序列。

穷举的不可行

要找全局最优序列，理论上需要枚举所有可能的 token 组合。假设词表大小为 V、序列长度为 T，候选序列数是 V^T（V 的 T 次方），呈指数爆炸。即便 V 取 3 万、T 取 30，组合数约为 3 万^30 ≈ 10^134，远远超过宇宙原子数（约 10^80），完全不可计算。

穷举搜索复杂度: O(V^T × T)  -- V 词表大小, T 序列长度
Beam Search 复杂度: O(k × V × T)  -- k 束宽, 每步扩展 k 条候选, 每条 V 个选择

当 k=5、V=30000、T=50 时，Beam Search 只需约 750 万次概率计算，而穷举需要 30000^50 次——差距超过 100 个数量级。

Beam Search 的核心机制

Beam Search 维护一个大小为 beam_width（束宽，记作 k）的候选序列集合，称为 beam（束）。算法流程如下：

初始化：beam 中只有一个起始 token（如 BOS，Beginning of Sequence，序列起始符）。
逐步扩展：在第 t 步，对 beam 中的每一个候选序列，用模型预测下一步的概率分布，扩展出词表中所有可能的下一个 token，每个扩展得到一条新候选。每步最多产生 k × V 条新候选。
打分：对新候选序列打分。分数通常用对数概率累加：score = log p(y_1) + log p(y_2) + … + log p(y_t)。之所以用对数，一是把连乘变成连加、避免数值下溢（多个小于 1 的概率相乘会迅速趋近 0），二是单调性保证”对数概率之和最大”等价于”原始概率之积最大”。
剪枝保留：把所有新候选按 score 降序排列，只保留前 k 条，其余丢弃，构成新的 beam。
终止：当某条候选生成结束符（EOS，End of Sequence，序列结束符）时，它从活跃 beam 中移出，进入”已完成序列”集合；同时 beam 中补入一条新的活跃候选以维持 k 条。当所有候选都结束（或达到最大长度）时，从已完成集合里选 score 最高的输出。

为什么用对数概率

原始联合概率 P = p_1 × p_2 × … × p_t，当 t 较大、每个 p_i 都小于 1 时，乘积会迅速下溢成浮点 0，丢失排序信息。取对数后 log P = log p_1 + … + log p_t，所有项都是负数相加，数值范围稳定得多。由于 log 是单调递增函数，比较 log P 的大小关系与比较 P 完全一致，不影响选优。

数值下溢示例（float64 最小正数约 4.9e-324）:

p_i = 0.1（每步概率，在词表较大的模型中很常见）

连乘（原始概率）:
  10 步:  0.1^10  = 1e-10     ← 还能表示
  50 步:  0.1^50  = 1e-50     ← 接近下溢边界
  200 步: 0.1^200 = 1e-200    ← 浮点数下溢为 0.0！

对数累加:
  10 步:  10 × log(0.1) = -23.03   ← 完全没有精度问题
  200 步: 200 × log(0.1) = -460.52  ← 依然稳定可比较

束宽（beam width）的选择

width = 1：退化为贪心解码，只保留一条最优，速度最快、质量最差。
width = 5 到 10：最常见的工程取值，质量提升明显、计算开销可控（约为贪心的 k 倍）。
width 越大：越逼近穷举的最优解，但算力与显存线性增长，且在开放式生成中收益递减——超过某个阈值后，更宽的 beam 往往只是生成更”安全平庸”的文本，而非更”好”的文本。

长度归一化（Length Normalization）

一个关键陷阱：对数概率是负数累加，序列越长，累加的负数越多，总 score 天然越低。如果直接比较原始 score，模型会偏好短序列（因为短序列”负得少”），导致输出过早截断。解决方法是对 score 按长度归一化，常用做法是除以长度（或长度的某个幂次 alpha）：

\text{normalized\_score} = \frac{\text{score}}{\text{length}^{\alpha}}

其中：

$\text{score} = \sum \log p(y_t \mid y_{1:t-1})$ — 原始累计对数概率
$\text{length}$ = 序列长度（已生成的 token 数）
$\alpha$ = 长度惩罚系数，通常取 0.6 到 1.0

alpha 的直觉理解：

alpha = 0：不做归一化，偏好短序列。
alpha = 0.6 到 0.8：最常见的工程取值（GNMT 用的是基于覆盖率的变体）。
alpha = 1.0：完全按”平均每步对数概率”比较，对长短序列一视同仁。
alpha > 1.0：鼓励更长的输出。

Google 在 2016 年的神经机器翻译工作（GNMT）中提出的长度惩罚（length penalty）就是这一思路的经典实现，其公式略有不同：

\text{lp} = \left(\frac{5 + \text{length}}{5 + 1}\right)^{\alpha}

\text{normalized\_score} = \frac{\text{raw\_score}}{\text{lp}}

其中常数 5 是经验值，用于在短序列时减少过度惩罚； $\alpha$ 通常取 0.6 到 1.0。

Diverse Beam Search（多样性束搜索）

标准 Beam Search 的 k 条候选往往高度相似——只在个别词上不同，本质上是同一条路径的微小变体。Diverse Beam Search（Vijayakumar et al., 2016）把 beam 分成 G 组，每组 k/G 条。组内正常打分，同时对与前面各组已选 token 的相似度施加惩罚：

\text{score\_diverse}(y) = \text{score\_logprob}(y) - \lambda \times \text{similarity\_penalty}(y, \text{previous\_groups})

其中：

$\lambda$ — 多样性惩罚强度（超参数）
$\text{similarity\_penalty}$ — 与前面各组中 token 的重叠程度

直觉上，这相当于告诉后面的组：“前面的人已经选过的路径，你要走就得额外扣分”，强制不同组走出不同的分支。这在对话、图像描述等需要”多选题式”候选的场景中很有用。

Beam Search vs 采样

维度	Beam Search	采样（Top-K / Top-P）
输出	确定性，同一输入永远同一输出	随机，每次可能不同
多样性	低，容易”正确但无聊”	高，富有变化
典型场景	翻译、摘要、结构化抽取（有标准答案）	创意写作、对话、头脑风暴
计算量	k 倍贪心	与贪心相当

经验法则：任务有明确对错时用 Beam Search，任务鼓励发散时用采样。 现代开源 LLM（如 Llama 系列）默认采样，而经典翻译模型（GNMT、早期 Transformer）默认 Beam Search。

原理图解

上图中，束宽 k=2：第一步从所有扩展中保留 A、B 两条，淘汰 C；第二步各自再扩展并保留前 2，最后从已完成序列里选总分最高者输出。

完整打分流程示意

以束宽 k=2、词表 {“猫”, “狗”, “鸟”} 为例，展示两步完整扩展：

第 2 步从 6 条候选中保留 score 最高的 2 条（“猫睡”和”猫追”），淘汰其余。

动手实践

1. numpy 手写 Beam Search（完整注释版）

下面用 numpy 手写一个最小化的 Beam Search，演示”扩展—打分—剪枝”的核心循环：

import numpy as np

def beam_search(predict_next, vocab, beam_width=3, max_len=10,
                eos_id=None, length_alpha=0.7):
    """
    手写 beam search，带长度归一化。

    参数:
        predict_next : 函数，输入已生成的 token id 列表，返回下一步概率数组
        vocab        : 词表列表，vocab[i] 是 id=i 对应的 token 字符串
        beam_width   : 束宽 k，每步保留的候选数
        max_len      : 最大生成长度
        eos_id       : 序列结束符的 id（可选）
        length_alpha : 长度归一化指数，0 到 1 之间
    """
    # 每个 beam 是一个元组：(累计对数概率, token id 列表, 是否已完成)
    beams = [(0.0, [<BOS_id>], False)]
    finished = []

    for step in range(max_len):
        candidates = []

        for score, tokens, done in beams:
            if done:  # 已完成的 beam 不再扩展
                candidates.append((score, tokens, True))
                continue

            # 模型预测下一步概率分布
            probs = predict_next(tokens)  # 返回 shape=(vocab_size,) 的概率数组

            # 取概率最高的 beam_width 个 token（减少候选数量）
            top_ids = np.argsort(probs)[::-1][:beam_width]

            for idx in top_ids:
                p = probs[idx]
                # 对数概率累加（加 1e-12 防止 log(0)）
                new_score = score + np.log(p + 1e-12)
                new_tokens = tokens + [idx]
                is_done = (idx == eos_id) if eos_id is not None else False
                candidates.append((new_score, new_tokens, is_done))

        # 按原始分数排序（长度归一化在最终输出时才做）
        candidates.sort(key=lambda x: x[0], reverse=True)

        # 分离已完成和活跃候选
        new_beams = []
        for score, tokens, done in candidates:
            if done:
                finished.append((score, tokens))  # 移入已完成集合
            else:
                new_beams.append((score, tokens, done))
            if len(new_beams) >= beam_width:
                break

        # 如果活跃 beam 不足，从已完成中"借"（补位以维持 beam 大小）
        beams = new_beams[:beam_width]
        if not beams:  # 所有 beam 都已完成
            break

    # 合并活跃和已完成的候选，做长度归一化后选最优
    all_candidates = finished + [(s, t) for s, t, _ in beams]

    def length_normalized(item):
        score, tokens = item
        length = len(tokens)
        return score / (length ** length_alpha)  # 长度归一化打分

    best = max(all_candidates, key=length_normalized)
    return [vocab[i] for i in best[1]]  # 返回 token 字符串列表

2. 用 mock 模型验证 Beam Search 行为

# 构造一个简单的 mock 模型来验证 Beam Search 的正确性
# 模拟一个只有 4 个 token 的词表
vocab = ["<BOS>", "猫", "追", "睡", "<EOS>"]
BOS_ID, EOS_ID = 0, 4

# 预定义的转移概率（mock predict_next）
# 实际中这一步由神经网络计算
transitions = {
    (0,):   [0.0, 0.6, 0.3, 0.1, 0.0],  # BOS → 猫(0.6) 追(0.3) 睡(0.1)
    (0, 1): [0.0, 0.0, 0.7, 0.3, 0.0],  # 猫 → 追(0.7) 睡(0.3)
    (0, 2): [0.0, 0.5, 0.0, 0.2, 0.3],  # 追 → 猫(0.5) 睡(0.2) EOS(0.3)
    (0, 1, 2): [0.0, 0.0, 0.0, 0.6, 0.4],  # 猫追 → 睡(0.6) EOS(0.4)
    (0, 1, 3): [0.0, 0.0, 0.1, 0.0, 0.9],  # 猫睡 → EOS(0.9)
}

def mock_predict(tokens):
    key = tuple(tokens)
    return np.array(transitions.get(key, [0.0, 0.2, 0.2, 0.2, 0.4]))

# 用束宽 2 运行
result = beam_search(mock_predict, vocab, beam_width=2, max_len=5, eos_id=EOS_ID)
print(" → ".join(result))
# 预期输出类似：<BOS> → 猫 → 追 → <EOS>

3. Hugging Face transformers 实际调用

工程实践中通常直接用 Hugging Face transformers 内置实现：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")
tok = AutoTokenizer.from_pretrained("gpt2")
inputs = tok("The future of AI is", return_tensors="pt")

# --- 基础 Beam Search ---
out = model.generate(
    **inputs,
    num_beams=5,              # 束宽 5
    max_new_tokens=30,
    length_penalty=0.8,       # 大于 0 鼓励更长输出
    no_repeat_ngram_size=2,   # 禁止 2-gram 重复，缓解啰嗦
    early_stopping=True,      # 所有 beam 都无法超越当前最优时提前停止
)
print(tok.decode(out[0], skip_special_tokens=True))

# --- Diverse Beam Search（多样性束搜索）---
out_diverse = model.generate(
    **inputs,
    num_beams=6,              # 总束宽
    num_beam_groups=3,       # 分 3 组，每组 2 条
    diversity_penalty=1.0,   # 组间多样性惩罚强度
    max_new_tokens=30,
    early_stopping=True,
)

# --- 返回多候选（用于排序/重排） ---
out_multi = model.generate(
    **inputs,
    num_beams=5,
    num_return_sequences=3,   # 返回 score 最高的 3 条候选
    max_new_tokens=30,
    early_stopping=True,
    return_dict_in_generate=True,
    output_scores=True,
)
for i, seq in enumerate(out_multi.sequences):
    print(f"候选 {i} (score={out_multi.sequences_scores[i]:.3f}): "
          f"{tok.decode(seq, skip_special_tokens=True)}")

4. Beam Search 打分可视化

import numpy as np

def visualize_beam_scores(scores_by_step):
    """
    可视化 beam search 每步的候选分数。
    scores_by_step: 列表的列表，每个子列表是该步所有候选的 score
    """
    for step, scores in enumerate(scores_by_step):
        print(f"步骤 {step}:")
        for i, s in enumerate(sorted(scores, reverse=True)):
            bar = "█" * int((-s) * 20)  # score 是负数，取绝对值画条
            marker = " ✓ 保留" if i < 2 else " ✗ 淘汰"  # 假设束宽 2
            print(f"  候选{i}: {s:8.3f} {bar}{marker}")
        print()

# 示例数据
visualize_beam_scores([
    [-0.4, -0.7, -2.1],           # 步骤 0: 保留前 2
    [-0.9, -0.7, -1.2, -1.3, -1.5],  # 步骤 1: 保留前 2
])

实践要点

束宽不是越大越好：从 1 到 5 质量提升明显，从 5 到 20 收益骤减，且推理时延线性上升；大多数任务 4 到 8 就够。
务必做长度归一化：不归一化的原始对数概率会强烈偏向短序列，是新手最常踩的坑。
配合 no_repeat_ngram_size：Beam Search 容易在局部打转、产生重复短语，加上 n-gram 重复惩罚（n-gram repetition penalty，即禁止输出中出现重复的 n 个连续 token）能显著缓解。
开放式生成慎用：在聊天、创意写作等无标准答案的场景，Beam Search 输出常被评价为”安全但乏味”，此时改用 Top-P 采样（核采样，Nucleus Sampling，只从累积概率超过 p 的最小 token 集合中采样）更合适。
注意 EOS 处理：候选生成 EOS 后应移出活跃 beam 并补位，否则 beam 会越走越少。
batch 推理要处理不等长：同一 batch 内不同序列完成步数不同，需要 padding（填充）与 attention mask（注意力掩码），这是实现层的复杂度大头。
可用 early_stopping：当 beam 中所有候选都已无法超过当前最优已完成序列时即可提前停止，节省算力。
约束式 Beam Search（Constrained Beam Search）：当需要输出满足特定格式（如 JSON Schema、正则表达式）时，可以在 Beam Search 的每步剪枝中只保留满足约束的 token。Hugging Face 的 model.generate(constraints=[...]) 和 Outlines / Guidance 等库都提供了这一能力，是 2024 年以来结构化输出的主流方案。

典型应用

机器翻译：GNMT、Marian、早期 Transformer 翻译模型几乎清一色用 Beam Search（束宽 4 到 5），是 BLEU 分数（BLEU Score，一种机器翻译质量评估指标）提升的标准武器。
文本摘要：CNN/DailyMail、XSum 等摘要任务中，Beam Search 配合长度惩罚是默认配置。
图像描述（Image Captioning）：Show and Tell 及其后继模型用 Beam Search 从图像特征生成描述候选。
语音识别：CTC + Beam Search decoding（如 ESPnet、OpenFST）是端到端 ASR（自动语音识别）的经典后处理流程。CTC（Connectionist Temporal Classification，联结时间分类）输出的是帧级概率分布，Beam Search 负责合并重复 token、处理空白符、搜索最可能的文字序列。
拼写纠错与 OCR 后处理：在候选字符序列上用 Beam Search 选最优校正路径。
代码补全（早期）：早期代码模型（如 CodeGPT）在补全场景用 Beam Search 提升建议质量；现代大模型更倾向采样。
结构化生成：2024 年以来，约束式 Beam Search 广泛用于让 LLM 输出合规的 JSON、SQL、函数调用等结构化格式，兼顾正确性和概率最优性。

典型类库与工具

类库	语言	说明
Hugging Face transformers	Python	`generate(num_beams=k)` 一行启用，支持长度惩罚、n-gram 约束、diverse beam、约束式生成
Fairseq	Python	Meta 的序列建模工具箱，Beam Search 实现成熟，翻译任务常用
OpenNMT	Python / Lua	开源神经机器翻译框架，Beam Search 与 length normalization 齐全
ESPnet	Python	语音与翻译端到端工具，CTC + Beam Search decoding
OpenFST / pyctcdecode	C++ / Python	加权有限状态机（Weighted Finite-State Transducer）与 CTC beam decoding，语音识别后处理
Outlines / Guidance	Python	约束式/引导式生成库，用文法或 JSON Schema 约束 Beam Search 输出格式
tensor2tensor	Python	Google 早期序列模型库，包含 length penalty 的经典实现

术语表

术语	英文	解释
束搜索	Beam Search	每步保留固定宽度个最优候选的序列解码算法
束宽	beam width	每步保留的候选数量，记作 k；k=1 即贪心
贪心解码	Greedy Decoding	每步取概率最大 token，不回头
穷举搜索	Exhaustive / Brute-force Search	枚举所有序列，计算量随长度指数爆炸
对数概率	Log Probability	概率取对数，把连乘变连加，防止下溢
分数 / 打分	Score	beam 候选的评价值，通常用累计对数概率
剪枝	Pruning	保留前 k 条候选、丢弃其余的过程
长度归一化	Length Normalization	对累计分数按序列长度修正，消除对短序列的偏好
长度惩罚	Length Penalty	长度归一化的具体形式，常用 score 除以 length 的 alpha 次幂
多样性束搜索	Diverse Beam Search	分组并施加组间多样性惩罚的 Beam Search 变体
约束式束搜索	Constrained Beam Search	在剪枝时加入格式约束（如 JSON Schema），确保输出满足结构化要求
结束符	EOS（End of Sequence）	标记序列生成结束的特殊 token
自回归生成	Autoregressive Generation	每步基于已生成内容预测下一 token 的生成方式
N-gram 重复惩罚	no_repeat_ngram_size	禁止输出中出现重复的 n 个连续 token，缓解 Beam Search 的重复输出问题
推测解码	Speculative Decoding	用小模型猜多步再由大模型验证，加速自回归推理