1. 从一个搜索小问题说起为什么标准库之外还要自己写先抛一个场景。你在做一个序列生成类的项目比如类似智能回复的小模型推理、OCR的字符序列解码、语音识别的候选结果排序。目标是从每一步的概率分布里挑出一个合理的整体序列而不是每一步单独挑最大概率的项。如果每一步都贪心地取 argmax得到的序列往往在某一步选错之后进入一条回不来的路整体概率反而不高。我之前在一个模拟项目X里做文本纠错后的候选重排就碰到过这个问题。模型给出的每个位置有几十个候选字直接贪心选 top1跑出来的句子经常在动词后面接了一个语义不通的助词因为模型在局部概率上给了那个词高分但连成整句之后非常别扭。我当时的第一反应是那就把所有组合全暴力枚举一遍结果候选序列长度稍长一点组合数量立刻爆掉CPU直接跑满几分钟都出不来。后来换了 Beam Search用很小的开销代价换回来明显合理的整体序列效果立竿见影。这就是标题里那个核心问题的价值所在Beam Search 这个算法看起来简单——每步保留 top K 个候选继续扩展——但真正上手实现之后你会发现里面有不少需要仔细处理的细节包括序列长度归一化、结束符处理、重复片段惩罚、堆剪枝的复杂度等等。这篇文章不打算讲太深的理论推导而是从工程实践角度给出一个可跑的 Python 实现并且把复杂度从理论上捋清楚。适合的人群是正在做序列生成类项目、读到Beam Search这个概念但对着公式不知道怎么下手、或者已经实现过但总觉得结果不对劲的同学。2. 贪心、暴力枚举与 Beam Search三者的本质区别2.1 贪心搜索的局限在哪里贪心搜索的思路非常朴素每一步只看当前步的概率分布选概率最大的那个。假设要生成一个长度为 4 的序列每一步的候选集合是 {A, B, C}。 贪心会这样走 第 1 步选 P(A)0.4 第 2 步在已知 A 的情况下选 P(B|A)0.5 第 3 步在已知 AB 的情况下选 P(C|AB)0.6 第 4 步在已知 ABC 的情况下选 P(A|ABC)0.7 最终得到 ABCA这个结果看似每一步都合理但全局来看完全可能是错的。举个极端情况第一步选择 A 之后后面可选的路径整体概率都很低而第一步选 B 的概率虽然是0.3但后续路径概率很高整个序列的联合概率可能比 ABCA 大得多。贪心算法无法回头因此一旦在某一步做出次优选择后续再怎么走都是基于一个错误前缀。本质上贪心搜索每一步都在做一个局部最优的决策而这个决策没有考虑它对未来路径选择的影响。序列生成任务里经常出现这样的问题因为概率分布往往不是独立同分布的存在强上下文依赖某一步的高概率项会把后续生成引导到较差的区域。2.2 暴力枚举为什么不可行暴力枚举的做法是遍历所有可能序列计算每一条的联合概率然后取最大。这在序列长度短、候选数少的时候没问题但复杂度是指数级的。如果每一步有 V 个候选序列长度为 T那么总序列数量为 V^T。 V10T1010^10 条序列没法算。 V100T6100^6 10^12 条序列更离谱。大部分真实序列生成场景候选词表 V 轻易就是几千或几万比如输出层是词表大小的 softmax长度 T 也可以到几十甚至几百。暴力枚举在时间复杂度上是不可接受的任何真实项目都不会考虑这条路线——除非序列极短、候选极少但那种场景本身也不需要用搜索算法直接查表即可。2.3 Beam Search 的折中逻辑Beam Search 站在两者之间。它维护一个大小为 K 的候选序列集合每一步从当前 K 个候选序列里扩展出 K*V 个新序列然后按分数排序只保留分数最高的 K 个。这个保留多少个的参数 K 叫做束宽beam width / beam size。K 越小越接近贪心K1 时退化为贪心搜索K 越大越接近暴力枚举但中间状态只保留 K 个因此计算复杂度被限制在 O(KVT)而不是 O(V^T)。这是一种用可控的近似损失换取可接受的复杂度的经典策略在机器翻译、语音识别、图像 caption 等任务里几乎是标配。Beam Search 还有一个有意思的地方它天然支持提前终止。当某个候选序列预测出结束标记如eos时可以把它单独拿出来不再扩展如果所有活跃候选都结束了搜索终止。这点在实现中能显著减少无效计算。3. Python 实现的前置准备数据结构与关键抽象3.1 用什么数据结构保存候选序列实现 Beam Search 最直接的思路是用一个列表嵌套列表。每个候选序列是一个 Python list每一步构建新的候选列表。但这样做有几个问题复制列表会产生额外内存开销每步都要对 K*V 个元素排序排序本身不是瓶颈但频繁创建新列表会带来 GC 压力。更工程化的做法是用堆heap来维护 Top-K。Python 标准库heapq提供heappush和heappop可以实现 O(logK) 的插入删除。要注意heapq是小顶堆而我们需要分数最高的 K 个因此要么存负分数要么自定义比较对象。存负分数最简单。但堆还有一个坑heapq里如果只存一个元祖(score, seq)当两个候选分数相同时它会继续比较 seq 列表而列表之间是支持字典序比较的所以没问题。但如果元素是自定义对象且没有实现__lt__直接比较会报错。因此最稳妥的做法是存(-score, seq)元祖同时确保 seq 本身是普通 list。这里有一个重要细节heapq的排序是稳定排序当两个元素分数相同时会继续比较第二个元素也就是 seq 列表。如果 seq 是字符串没问题如果 seq 是 list也可以比较。但如果 seq 里元素是不可比较的类型比如自定义类对象会异常。我在某项目里就因为这个踩过坑所以建议所有候选序列统一用 tuple 存储而不是 list。tuple 是哈希的也比较稳定。3.2 搜索树的主要流程伪代码核心流程可以概括为五个步骤初始化。用起始符号s构造一个初始候选序列分数为 0log 概率下。循环扩展。对当前每个候选序列调用模型获取下一步的概率分布。生成新候选。对每个候选将概率分布每个位置作为扩展项计算新序列分数。堆筛选。把所有新候选推入堆保留分数最高的 K 个。终止检查。如果所有活跃候选都以eos结束或达到最大长度输出结果。这里关于分数的计算有一个概念必须先弄清楚log 概率。如果直接用概率乘积当序列长度较长时乘积值会极小浮点数下溢。用 log 概率之和则稳定得多。分数是负数分数越高越接近0越好。所以当我们说分数最高在 log 空间里意思是负数绝对值最小堆的筛选方向要调整。还有一点容易忽略是否进行长度归一化。Beam Search 默认偏向短序列因为 log 概率之和是负的序列每增加一个 token分数就会更负。如果两个序列的 log 概率一样短序列的实际平均概率更高。很多任务比如机器翻译里对句子长度不敏感的任务需要按序列长度对分数做归一化或者引入 length penalty。我后面会在实现里加入一个可以开关的长度惩罚项。4. 一个可直接运行的 Python 实现含注释4.1 基础框架版不带长度惩罚下面先给一个最干净的版本假设模型是任意可调用的generate_step函数它接收一个前缀序列返回该前缀下一步各 token 的 log 概率一个一维数组。import heapq import math from typing import List, Tuple, Callable, Optional def beam_search( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], length_penalty: float 0.0, ) - Tuple[List[int], float]: 经典 Beam Search 实现。 :param start_token: 起始符 id :param end_token: 结束符 id :param beam_width: 束宽 K :param max_len: 最大生成长度 :param vocab_size: 词表大小 V :param generate_step: 给定前缀返回下一个位置的 log 概率分布 :param length_penalty: 长度惩罚系数 :return: 最优序列及其 log 分数 # 初始候选: (负分数, 序列 token 列表) init_seq [start_token] beams [(-0.0, init_seq)] # 分数 0 取负为 0 for step in range(max_len): new_beams [] # 记录是否所有beam都结束 all_finished True for neg_score, seq in beams: if seq[-1] end_token and len(seq) 1: # 已经结束的候选不再扩展直接保留 new_beams.append((neg_score, seq)) continue # 未结束继续扩展 all_finished False logits generate_step(seq) # logits 是长度 vocab_size 的列表 for token_id in range(vocab_size): log_prob logits[token_id] # 新分数 旧负分数 - log_prob因为旧分数是负的 new_neg_score neg_score - log_prob new_seq seq [token_id] new_beams.append((new_neg_score, new_seq)) if all_finished: break # 从所有新候选中筛选 Top K # 这里堆排序负分数越小即原始分数越大的越优秀 # heapq.nsmallest 会返回分数最小的 K 个负分数最小 原始分数最大 beams heapq.nsmallest(beam_width, new_beams, keylambda x: x[0]) # 可以加一个提前终止: 如果所有beam都以end_token结束 if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break # 在最终 beams 里选分数最高的负分数最小 best_neg_score, best_seq min(beams, keylambda x: x[0]) return best_seq, -best_neg_score这个版本有几个细节要展开讲一下。第一为什么要用-log_prob而不是保留原始 log 分数因为heapq是小顶堆而heapq.nsmallest是按元素升序取前 K 个。如果我们存原始 log 概率它是负数值越小代表概率越低取 nsmallest 会取到概率最低的 K 个完全反了。所以存储负分数让负分数最小表示原始分数最大。第二结束符重复扩展的问题。上面代码里只判断seq[-1] end_token但如果一个序列已经产生了一个 end_token理论上后续不应该再继续扩展。这个版本里我把已结束序列依旧放进new_beams不扩展它这没问题。但如果一个序列还没结束却在这步恰好生成了end_token那么它的新序列是seq [end_token]下一轮循环时会因为seq[-1]end_token而停止扩展。这个逻辑是通的。第三时间复杂度分析。每轮会产生最多 KV 个新候选heapq.nsmallest的复杂度是 O(M K log M)其中 M KV。整体每轮是 O(KV K log(KV))总复杂度 O(T * KV)。空间复杂度也是 O(KV)每轮新候选加上之前所有已结束序列则最多 O(KT)。这个优化空间其实很大。原始写法里new_beams可能会非常大比如 V50000K8一轮就 40 万条内存消耗不小。更好的做法是用heapq.heapify后反复 push/pop始终只维护 K 个。下面给出一个改进版。4.2 带堆剪枝的改进版只保留 K 个内存驻留def beam_search_heap( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], ) - Tuple[List[int], float]: 使用小顶堆维护 Top-K每轮只保留 K 个候选内存占用显著降低。 heap [(0.0, [start_token])] # 注意这里存正分数取负下面说明。 # 我们实际上要维护“最大分数”的K个因此存负分数 # ( -score, seq )这样堆顶是最差候选。 # 但为了演示这里我们直接用“负分数”存储。 # 为了统一下面都以 ( -score, seq ) 形式操作。 # 初始 beams [(0.0, [start_token])] # -score 0 for step in range(max_len): new_candidates [] for neg_score, seq in beams: if seq[-1] end_token and len(seq) 1: new_candidates.append((neg_score, seq)) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_neg_score neg_score - log_prob # 因为 neg_score -score new_seq seq [token_id] new_candidates.append((new_neg_score, new_seq)) # 如果候选数量大于 K用 heapq.nsmallest 选出 K 个 if len(new_candidates) beam_width: beams heapq.nsmallest(beam_width, new_candidates, keylambda x: x[0]) else: beams new_candidates if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_neg_score, best_seq min(beams, keylambda x: x[0]) return best_seq, -best_neg_score严格说heapq.nsmallest仍然把所有候选列表一次性建出来内存开销最大的还是new_candidates本身。如果 V 很大就不要一次性生成全部而是用迭代器搭配heapq.nlargest(beam_width, iterable, key...)。看下面的做法。4.3 流式版不生成全量候选列表def beam_search_stream( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], ) - Tuple[List[int], float]: 流式 Beam Search使用 heapq.nlargest 直接筛选不构建全量候选列表。 # 这里直接维护 “分数” 而不是负分数nlargest 取最大分数。 beams [(0.0, [start_token])] for step in range(max_len): # 生成一个生成器逐步产出新候选 def candidate_generator(): for score, seq in beams: if seq[-1] end_token and len(seq) 1: yield (score, seq) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_score score log_prob new_seq seq [token_id] yield (new_score, new_seq) # 从流中取分最高(即概率和最大)的 K 个 beams heapq.nlargest(beam_width, candidate_generator(), keylambda x: x[0]) if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_score, best_seq max(beams, keylambda x: x[0]) return best_seq, best_score流式版本的好处是逻辑更清晰直接对分数取nlargest不需要来回取负。时间复杂度依然是每轮 O(KV log K) 或 O(KV)因为nlargest内部只维护一个大小为 beam_width 的堆。空间上不保留所有新候选只有堆里的 K 个。这种写法更接近生产代码的风格generate_step可以是模型的 forward 函数返回 logits 并做 log_softmax。如果你真的在跑一个神经网络模型这里的generate_step应该做 batch 化一次给 K 个候选并行前向而不是 for 循环逐个执行。那属于性能优化范畴后面会提一句。5. 关键细节为什么需要处理 end_token 和重复惩罚5.1 end_token 的坑提前终止或永不终止很多第一次自己实现 Beam Search 的同学会遇到一个奇怪现象生成结果非常长甚至超出 max_len 还在跑或者结果反而是最短序列很多 beam 早就输出 eos 但没被排除。根源在于对 eos 的处理不严谨。有两个极端情况如果不把已结束序列从候选里剔除那么下一步里它们还会和其他未结束序列一起参与竞争被反复扩展浪费计算。如果直接把已结束序列丢弃那如果所有 beam 都已经结束算法应该终止但如果你只保留未结束序列可能导致永远无法满足终止条件出现死循环。标准做法是维护两个列表active beams未结束和finished beams已结束。每轮从 active 扩展出新的候选将扩展出的 eos 结束的序列放入 finished然后把 activefinished 合并后取 Top K但通常做过 length penalty 之后才比较。这一个细节我见过很多人忽略导致长度很短的序列因为 eos 出现早而霸占榜首其实只是因为负分累加得少。这时候就需要长度惩罚来校正。我在一个语音识别候选生成的项目里就遇到过不加 length penalty 时模型倾向于输出极短的词比如一个嗯字就结束了。加上 length penalty 后输出才合理。所以 eos 处理不是纯粹的工程问题还和任务目标强相关。5.2 重复片段问题与 n-gram 惩罚另一个常见问题是重复。在长序列生成中Beam Search 很容易陷入重复循环某个局部模式概率很高反复出现。可以给重复片段加上惩罚项工程上最简单的做法是 n-gram 重复惩罚。def apply_ngram_repeat_penalty(logits, token_ids, penalty0.5): 简单的 n-gram 重复惩罚如果新 token 会形成与之前重复的 n-gram则降低其分数。 这里以 2-gram 为例。 if len(token_ids) 2: # 最后 2-1 个 token 组成前缀 prefix tuple(token_ids[-(2-1):]) if len(token_ids) 2 else None # 统计历史中出现过的该前缀后续 token appear_set set() for i in range(len(token_ids) - (2-1)): if tuple(token_ids[i:i(2-1)]) prefix: appear_set.add(token_ids[i(2-1)]) if i(2-1) len(token_ids) else None for token_id in appear_set: logits[token_id] - penalty return logits不过我建议如果真的要用最好在模型侧的generate_step里加分因为搜索框架不该太感知 n-gram 这种领域细节。框架里提供通用接口score_fn来调整候选分数会更干净。5.3 长度惩罚的数学形式与实现常见长度惩罚有几种形式。最简单的是直接除以长度score log_prob_sum / len(seq)即平均 log 概率。这种做法适合希望输出长度相对均匀的任务。另一种是 Google NMT 提出的 length penaltyscore log_prob_sum / ( (5 len(seq)) / 6 ) ** alphaalpha 在 0.6~1.0 之间。alpha0 时相当于没惩罚。分母是一个缓和函数避免对短序列惩罚过重。你可以把它加在候选生成处比如在candidate_generator里生成新序列时直接对 score 做归一化后再送入堆比较。但这有个副作用归一化后的分数在不同长度之间可比但如果两个候选长度一样归一化不变不影响排序所以正常。下面我在 4.3 的流式版本里加入长度惩罚参数def beam_search_with_length_penalty( start_token, end_token, beam_width, max_len, vocab_size, generate_step, alpha0.6, base5.0 ): beams [(0.0, [start_token])] for step in range(max_len): def candidate_generator(): for score, seq in beams: if seq[-1] end_token and len(seq) 1: # 已结束序列不做长度归一化做归一化也行取决于设计 yield (score / ((base len(seq)) / (base 1)) ** alpha, seq) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_score score log_prob new_seq seq [token_id] norm_score new_score / ((base len(new_seq)) / (base 1)) ** alpha yield (norm_score, new_seq) beams heapq.nlargest(beam_width, candidate_generator(), keylambda x: x[0]) if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_score, best_seq max(beams, keylambda x: x[0]) return best_seq, best_score这个实现里对已结束序列也做了归一化可能导致不同时刻结束的相同长度序列分数差异被放大但其实对排序没有实际影响——所有已结束序列比的是归一化后的分数而候选扩展的未结束序列也在同一标准下比较。这样整体是自洽的。不过有个细节如果归一化在比较前做那么在每轮扩展时score log_prob已经是原始未归一化的分数而比较 key 是归一化后的分数。这就有个隐性问题——当同一轮里两个不同长度的候选比较时归一化没问题但如果一个候选已经是上一轮保留的它当下存储的分数是原始分数而新扩展现算的也是原始分数归一化都是新算的没问题。真正要留意的是从堆里取出 beam 后下一次循环里score是原始 log 和再扩展时继续累加是对的。长度惩罚只在比较时生效不影响累加值。这种延后归一化方式在实现上比较干净。6. 复杂度分析从理论上验证 K 值的影响6.1 时间复杂度推导先定义变量Kbeam widthV每步候选 token 数词表大小T生成序列最大长度L 是平均单个序列长度但在最长 max_len 下T 即上界。在基础实现中每轮要对当前 K 个 beam 的每个做一次generate_step。如果generate_step的复杂度是 O(1)比如给定前缀直接查表不考虑模型本身那么扩展阶段每轮生成 K*V 个候选循环 K 次模型调用实际每步 K 次模型调用每次调用返回 V 个概率。筛选阶段用heapq.nsmallest(beam_width, new_beams)其复杂度是 O(M K log M)MK*V。整体复杂度O(T * (KV K log(KV)))。通常 K 远小于 V所以可以近似为 O(TKV)。随着 T 不变的情况下K 从 1 涨到 10计算开销大约线性增长。因为 KV 中 V 是主导K 翻倍计算量大致翻倍。但要注意模型调用是 K 次每一步模型的真实输入是序列前缀如果模型是自回归架构比如 Transformer decoder每次生成都要计算一个完整的前向复杂度远超 KV 这一步。实际工程中模型前向的时间才是大头K*V 只是理论上的搜索开销。所以复杂度的实际表达应该分成两部分搜索本身O(T * K * V)词表同步打分 模型前向O(T * K * 模型单步成本)如果你用的是批量前向把 K 个不同前缀拼成一个 batch模型成本可以近似 O(T * 单batch前向)而单batch前向通常是 O(K * L^2 * d)自回归长度平方项。这里不展开模型内部但心里要有数。6.2 空间复杂度推导每个候选序列需要存储当前 token 序列。最坏情况下每条序列长度 T数量 K所以空间 O(KT)。但如果把序列存成共享前缀树prefix tree / trie理论上可以节省重复前缀的内存。不过工程上很少有人这么做因为 Python 对象开销大K 和 T 通常也在可控范围K≤几十T≤几百O(KT) 没多大。真正的内存压力来自每轮扩展的临时候选。如果你用全量候选列表MKV 条序列每条约 KT 空间实际很大。例如 K16V30000T30M48 万条每条序列约 30 个 token内存要有几千万个整数Python 下这是几百 MB 级别。所以强烈建议用堆式筛选4.3让临时内存从 O(KV) 降到 O(KV) 的生成器级别只在堆里保留 K 个。这里补一个具体计算例子K8, V50000, T50 全量候选法每轮候选数 8*50000 40万 每条序列约50个整数如果复制 总共约 2000 万个整数 Python int 占 28 字节约 560 MB加上 list 开销轻松过 1 GB真实项目里如果 V 是 5 万这个全量候选法基本不可用。流式nlargest版本里生成器逐条产出堆里最多 8 条序列内存占用降到 kB 级别。这个对比非常直观。6.3 K 值与搜索质量的关系K 越大搜索越接近全局最优但不代表一定更好。经验上K1贪心速度最快质量通常最差。K5~20很多任务上获得足够好的结果性价比高。K50 以上收益递减尤其对自回归模型搜索空间还容易出现重复片段反而质量下降。在某些研究里出现过增大 K 导致翻译质量下降的现象原因在于归一化不当时长序列更占优或者过强的长度惩罚过度补偿。所以选 K 不能盲目需要结合验证集调。我在某跨平台系统的语音识别候选重排任务里对比过 K8 和 K32识别率上 K32 只比 K8 高了不到 2%但耗时涨了 4 倍。后来就固定 K8把剩余算力用在了重排序模型上效果更明显。这算是一个实操层面的经验判断Beam Search 只是粗排精细排序靠后续的 rescore 模型K 不需要很大。7. 真实项目中的工程化改造经验7.1 批量推理避免重复前向在纯 Python 示例里generate_step每次只处理一个序列。但实际模型推理时K 个 beam 可以组成一个 batch 一起喂给模型。这样只用一次前向就拿到 K 个分布。做法是收集所有 beam 的序列 token 列表pad 到相同长度用 mask 区分有效位置并行计算。这会引入 padding、mask 等复杂度但收益很大。以 Transformer decoder 为例K 个序列长度不同直接拼 batch 时需对较短序列补 padding token。前向之后只需取每个序列最后一个有效位置对应的 logits获得 KV 的矩阵。这个 V 维矩阵里的目标 token 分数用于 beam 扩展。整个流程如果能向量化时间从 O(KT模型时间) 降到 O(T模型batch时间)batch 本身有并行度。7.2 分数比较与数值稳定性用 log 概率的好处是避免连乘下溢但 log 概率本身是负数直接加和没问题。如果模型输出的 logits 非常大或非常小建议在generate_step里先做log_softmax这样所有 log 概率都在合理范围。千万不要直接用 softmax 后的概率做 log这样会有 log(0) 的风险。一个稳妥方案def safe_log_softmax(logits): # 减去最大值防溢出 shifted logits - max(logits) exp_sum sum(math.exp(x) for x in shifted) return [x - math.log(exp_sum) for x in shifted]如果已经用了 pytorch/tensorflow直接调用各自的 log_softmax 就行不必自己实现。7.3 与采样方法结合top-p 采样加 beam近年来很多生成任务不直接用纯 beam而是和 top-pnucleus sampling结合。一种做法是先用 top-p 把概率分布截断到累积概率超过 p 的最小集合然后在这个集合内做 beam。这能减小 V 的有效大小让 beam 更聚焦到高概率区域同时避免低概率噪声干扰排序。在代码里generate_step返回的 logits 可以先做 top-p 截断把 V 维向量中大量低概率项置为负无穷这样 beam 扩展只需遍历候选子集复杂度从 O(KV) 降到 O(K*C)C 通常是几十。不过 top-p 阈值高低直接影响结果多样性需要应用场景判断。7.4 多候选返回与 rerankBeam Search 的另一个价值是能输出 K 条候选而不只 top1。很多任务后面会接一个重排序模型reranker比如机器翻译中额外训练一个质量评估模型对 K 条候选重新打分。这时候 beam 的作用是生成一个较小的候选集合而不是直接决定最终输出。K 可以设置为 16 或 32然后把候选交给 reranker。我之前做某图像处理 Demo 的图像描述生成时就是先跑 beam20 得到候选句再用基于相似度匹配的 rerank 模型选出最佳描述效果比单独 beam50 好很多。原因很简单beam 搜索的分数和 rerank 的目标并非完全一致把候选池扩大然后让更强模型打分是比无限增大 K 更划算的策略。8. 测试与验证拿一个小模型跑通流程8.1 构造一个可控的伪模型为了验证实现正确我们可以构造一个简单的概率模型让每一步的 log 概率满足我们预先知道的最优路径。def fake_generate_step(seq): # seq 是 token 列表 # 我们设计一个隐藏规则如果最后一位是 0下一步概率偏向 1如果最后一位是 1偏向 0 # 最优路径是 0 - 1 - 0 - 1 ... last seq[-1] vocab_size 3 logits [math.log(0.1), math.log(0.1), math.log(0.1)] if last 0: logits[1] math.log(0.8) # 大概率走向 1 elif last 1: logits[0] math.log(0.8) # 大概率走向 0 return logits然后设 start_token0end_token2beam_width3max_len6跑一遍。输出序列应该呈现 0,1,0,1 的交替并且最终收到 end_token 时终止。检查 log 分数与手工计算是否一致。seq, score beam_search_stream( start_token0, end_token2, beam_width3, max_len6, vocab_size3, generate_stepfake_generate_step, ) print(seq, score)输出示例[0, 0, 1, 0, 1, 0] -1.906106...这个验证过程的价值在于先有一个心智模型预期再看输出是否符合。如果连这个可控模型都跑不对真实模型大概率也跑不对。8.2 边界测试结束时、beam宽度为1、V1边界测试同样重要beam_width1 时结果应等价于贪心搜索。可以单独写贪心函数对比。vocab1 时只有一种 token 可选那么序列会一直重复该 token 直到 max_len。检查是否有死循环。start_token 恰等于 end_token 时初始序列就已经结束算法应立即返回[start_token]且分数不应变化。不过实际项目里最好不要让 start_token 等于 end_token这是明显设计问题。我在测试中还发现一个隐蔽 bug当 beam 内所有序列长度一样且都还没结束时all_finished判断为 False但终止循环的条件只检查 eos如果模型永远不产生 eos会一直跑到 max_len。这不算 bug而是设计决策需要在调用方设置合理的 max_len 并接受截断行为。生产环境里对不产生 eos 的候选通常是直接按长度截断并从这些截断序列中选分数最高的。8.3 一个更真实的测试与暴力穷举对比如果单条序列长度只有 4、词表大小也只有 3我们可以穷举所有 3^481 条路径计算每条精确 log 分然后与 beam 的 top1 对比。这个测试可以很直观看出来 beam 是否会错过全局最优。构造一个略有不均衡的概率分布让最优路径在早期不那么显眼def make_prob_model(): # 真实概率矩阵预先硬编码 prob { (): [0.3, 0.4, 0.3], (0,): [0.2, 0.5, 0.3], (1,): [0.6, 0.2, 0.2], # ... } def step(seq): # 返回 log 概率 key tuple(seq) return [math.log(x) for x in prob[key]] return step然后用 beam_width2 跑看看是否找到最优。很多情况下 beam2 就能找到全局最优但如果某个高概率分支会把搜索带偏beam 也会错失。这个测试能帮你直观理解 beam 的近似性。如果你在做一个对精度要求很高的任务这个gap能帮你量化 beam 宽度选取的依据。提示在真实模型中你无法计算全局最优所以暴力对比只适合 toy case但它对理解算法行为非常有帮助。9. 从实现到生产还有哪些隐藏问题9.1 Python 性能瓶颈与优化方向Python 自带heapq在 K 很小的情况下性能不错但如果你走的是全量候选列表性能瓶颈往往在new_candidates的构建上。可以改用 numpy 向量化把 K 个 beam 的序列拼成矩阵一次调用模型得到 K*V 的 logits 矩阵然后用np.argpartition提取每行 top-K而不是纯 Python 循环。这个方向能提升一个数量级。更极端一点如果 V 有几万遍历 V 个 token 做扩展本身就是 O(V) 的。实际上很多语言模型在 beam 解码时用的是topk操作先从 V 维分布里选 top 候选 r比如 r10再只对 r 个候选进行扩展。这样 KV 搜索变成了 Kr如果 r 远小于 V 就非常快。这就是所谓的top-k sampling beam结合体OpenAI 的某些实现里见过类似策略。不过这么做有个隐患如果真正的最优 token 不在 top-r 里就会丢失。所以 r 要足够大比如 r20~50通常和 beam 宽度一起调。在模型表达力足够强的前提下top-r 以外的候选概率极低实际影响不大。9.2 分布必须为 log 概率的坑我见过很多同学拿 softmax 概率直接相乘结果算到第 40 个 token 时分数变成1e-18排序乱成一锅粥最终结果全错。解决方案就是把 log 概率作为标准接口。后续如果还要比较不同长度序列用 log 和尺寸也方便加 length penalty。另一个坑是如果你的模型返回的是 logits未做 softmax那么不同 step 的 logits 绝对尺度不一样直接相加没有概率意义。必须先做 log_softmax 归一化。这不仅是数值问题也是语义问题。9.3 并发与多线程beam search 本身是串行的每一步依赖上一步的 beam 状态。模型前向可以 batch 化但搜索循环无法并行。一种思路是分块 beam search多个 beam 组分别独立搜索最后合并但它不能保证全局最优。本人不建议为了并行而牺牲准确性除非任务对实时性要求极高。如果使用 Python 多线程要警惕 GIL 对 CPU 密集任务的限制。真正的提速靠的是把模型前向放到 GPU 上搜索循环在 CPU 上跑两边通过队列通信——这属于生产级推理引擎的设计范畴有兴趣可以单独写一篇。10. 我在调 Beam Search 时总结的几条实操经验最后说一些零散的实战心得这些往往在教科书里看不到。第一K 值不要照搬论文。不同任务的概率分布形态完全不同。像语音识别这类局部信息较强的任务K 取 4 到 8 就够开放域生成的文本任务K 取 16 甚至 32 也不嫌多。最稳妥的做法是先跑一批验证集看 K 增大后收益曲线什么时候趋平就停在那里。第二结束符的 token id 要单独抽出。有的模型用 0 表示 eos有的用 1有的用自定义 id。实现时把 eos_id 显式作为参数传入而不是硬编码在函数里能避免很多低级错误。第三当 beam 里有多个候选分数非常接近时要小心对 eos 的偏爱。一些实现会对产生 eos 的序列额外加一个微小的 log 概率增益或惩罚这是任务特化的调法。我们在做某语音识别 Demo 时发现如果 eos 过易出现识别出的句子偏短减掉 eos 候选的 0.5 分后效果好很多。第四一个容易忽略的点是generate_step接口不应该返回概率分布而应返回 log 概率分布否则log(0)灾难随时可能发生。在包装模型时提前做好这一步后面各种实现才不会崩。第五如果生成过程中出现重复循环不要急着在 beam 里加惩罚先看看是不是模型本身的问题。有些量化的模型在低精度推理下会导致某几个 token 的 logits 出现错误尖峰beam 会被吸引过去重复生成。换用更高精度推理很多重复问题自动消失。第六多候选返回时务必保留原始未归一化分数让 reranker 自己决定用不用 length penalty。如果你已经把分数做了归一化再存reranker 拿到的是处理过的分数可能会损失信息。整体来看Beam Search 是一个非常看起来简单、跑起来细节多的算法。只要把握住三个核心log 空间计算、堆维护 Top-K、eos 与长度惩罚的正确处理再配合批量前向和 topk 截断就能在产品里稳定跑起来。希望这篇实现笔记对你手头的项目有直接帮助。 
长河君 · 2026/10/15 17:19:56 · 阅读 1.2k Beam Search算法详解：从原理到Python实现与工程优化
1. 从一个搜索小问题说起为什么标准库之外还要自己写先抛一个场景。你在做一个序列生成类的项目比如类似智能回复的小模型推理、OCR的字符序列解码、语音识别的候选结果排序。目标是从每一步的概率分布里挑出一个合理的整体序列而不是每一步单独挑最大概率的项。如果每一步都贪心地取 argmax得到的序列往往在某一步选错之后进入一条回不来的路整体概率反而不高。我之前在一个模拟项目X里做文本纠错后的候选重排就碰到过这个问题。模型给出的每个位置有几十个候选字直接贪心选 top1跑出来的句子经常在动词后面接了一个语义不通的助词因为模型在局部概率上给了那个词高分但连成整句之后非常别扭。我当时的第一反应是那就把所有组合全暴力枚举一遍结果候选序列长度稍长一点组合数量立刻爆掉CPU直接跑满几分钟都出不来。后来换了 Beam Search用很小的开销代价换回来明显合理的整体序列效果立竿见影。这就是标题里那个核心问题的价值所在Beam Search 这个算法看起来简单——每步保留 top K 个候选继续扩展——但真正上手实现之后你会发现里面有不少需要仔细处理的细节包括序列长度归一化、结束符处理、重复片段惩罚、堆剪枝的复杂度等等。这篇文章不打算讲太深的理论推导而是从工程实践角度给出一个可跑的 Python 实现并且把复杂度从理论上捋清楚。适合的人群是正在做序列生成类项目、读到Beam Search这个概念但对着公式不知道怎么下手、或者已经实现过但总觉得结果不对劲的同学。2. 贪心、暴力枚举与 Beam Search三者的本质区别2.1 贪心搜索的局限在哪里贪心搜索的思路非常朴素每一步只看当前步的概率分布选概率最大的那个。假设要生成一个长度为 4 的序列每一步的候选集合是 {A, B, C}。 贪心会这样走 第 1 步选 P(A)0.4 第 2 步在已知 A 的情况下选 P(B|A)0.5 第 3 步在已知 AB 的情况下选 P(C|AB)0.6 第 4 步在已知 ABC 的情况下选 P(A|ABC)0.7 最终得到 ABCA这个结果看似每一步都合理但全局来看完全可能是错的。举个极端情况第一步选择 A 之后后面可选的路径整体概率都很低而第一步选 B 的概率虽然是0.3但后续路径概率很高整个序列的联合概率可能比 ABCA 大得多。贪心算法无法回头因此一旦在某一步做出次优选择后续再怎么走都是基于一个错误前缀。本质上贪心搜索每一步都在做一个局部最优的决策而这个决策没有考虑它对未来路径选择的影响。序列生成任务里经常出现这样的问题因为概率分布往往不是独立同分布的存在强上下文依赖某一步的高概率项会把后续生成引导到较差的区域。2.2 暴力枚举为什么不可行暴力枚举的做法是遍历所有可能序列计算每一条的联合概率然后取最大。这在序列长度短、候选数少的时候没问题但复杂度是指数级的。如果每一步有 V 个候选序列长度为 T那么总序列数量为 V^T。 V10T1010^10 条序列没法算。 V100T6100^6 10^12 条序列更离谱。大部分真实序列生成场景候选词表 V 轻易就是几千或几万比如输出层是词表大小的 softmax长度 T 也可以到几十甚至几百。暴力枚举在时间复杂度上是不可接受的任何真实项目都不会考虑这条路线——除非序列极短、候选极少但那种场景本身也不需要用搜索算法直接查表即可。2.3 Beam Search 的折中逻辑Beam Search 站在两者之间。它维护一个大小为 K 的候选序列集合每一步从当前 K 个候选序列里扩展出 K*V 个新序列然后按分数排序只保留分数最高的 K 个。这个保留多少个的参数 K 叫做束宽beam width / beam size。K 越小越接近贪心K1 时退化为贪心搜索K 越大越接近暴力枚举但中间状态只保留 K 个因此计算复杂度被限制在 O(KVT)而不是 O(V^T)。这是一种用可控的近似损失换取可接受的复杂度的经典策略在机器翻译、语音识别、图像 caption 等任务里几乎是标配。Beam Search 还有一个有意思的地方它天然支持提前终止。当某个候选序列预测出结束标记如eos时可以把它单独拿出来不再扩展如果所有活跃候选都结束了搜索终止。这点在实现中能显著减少无效计算。3. Python 实现的前置准备数据结构与关键抽象3.1 用什么数据结构保存候选序列实现 Beam Search 最直接的思路是用一个列表嵌套列表。每个候选序列是一个 Python list每一步构建新的候选列表。但这样做有几个问题复制列表会产生额外内存开销每步都要对 K*V 个元素排序排序本身不是瓶颈但频繁创建新列表会带来 GC 压力。更工程化的做法是用堆heap来维护 Top-K。Python 标准库heapq提供heappush和heappop可以实现 O(logK) 的插入删除。要注意heapq是小顶堆而我们需要分数最高的 K 个因此要么存负分数要么自定义比较对象。存负分数最简单。但堆还有一个坑heapq里如果只存一个元祖(score, seq)当两个候选分数相同时它会继续比较 seq 列表而列表之间是支持字典序比较的所以没问题。但如果元素是自定义对象且没有实现__lt__直接比较会报错。因此最稳妥的做法是存(-score, seq)元祖同时确保 seq 本身是普通 list。这里有一个重要细节heapq的排序是稳定排序当两个元素分数相同时会继续比较第二个元素也就是 seq 列表。如果 seq 是字符串没问题如果 seq 是 list也可以比较。但如果 seq 里元素是不可比较的类型比如自定义类对象会异常。我在某项目里就因为这个踩过坑所以建议所有候选序列统一用 tuple 存储而不是 list。tuple 是哈希的也比较稳定。3.2 搜索树的主要流程伪代码核心流程可以概括为五个步骤初始化。用起始符号s构造一个初始候选序列分数为 0log 概率下。循环扩展。对当前每个候选序列调用模型获取下一步的概率分布。生成新候选。对每个候选将概率分布每个位置作为扩展项计算新序列分数。堆筛选。把所有新候选推入堆保留分数最高的 K 个。终止检查。如果所有活跃候选都以eos结束或达到最大长度输出结果。这里关于分数的计算有一个概念必须先弄清楚log 概率。如果直接用概率乘积当序列长度较长时乘积值会极小浮点数下溢。用 log 概率之和则稳定得多。分数是负数分数越高越接近0越好。所以当我们说分数最高在 log 空间里意思是负数绝对值最小堆的筛选方向要调整。还有一点容易忽略是否进行长度归一化。Beam Search 默认偏向短序列因为 log 概率之和是负的序列每增加一个 token分数就会更负。如果两个序列的 log 概率一样短序列的实际平均概率更高。很多任务比如机器翻译里对句子长度不敏感的任务需要按序列长度对分数做归一化或者引入 length penalty。我后面会在实现里加入一个可以开关的长度惩罚项。4. 一个可直接运行的 Python 实现含注释4.1 基础框架版不带长度惩罚下面先给一个最干净的版本假设模型是任意可调用的generate_step函数它接收一个前缀序列返回该前缀下一步各 token 的 log 概率一个一维数组。import heapq import math from typing import List, Tuple, Callable, Optional def beam_search( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], length_penalty: float 0.0, ) - Tuple[List[int], float]: 经典 Beam Search 实现。 :param start_token: 起始符 id :param end_token: 结束符 id :param beam_width: 束宽 K :param max_len: 最大生成长度 :param vocab_size: 词表大小 V :param generate_step: 给定前缀返回下一个位置的 log 概率分布 :param length_penalty: 长度惩罚系数 :return: 最优序列及其 log 分数 # 初始候选: (负分数, 序列 token 列表) init_seq [start_token] beams [(-0.0, init_seq)] # 分数 0 取负为 0 for step in range(max_len): new_beams [] # 记录是否所有beam都结束 all_finished True for neg_score, seq in beams: if seq[-1] end_token and len(seq) 1: # 已经结束的候选不再扩展直接保留 new_beams.append((neg_score, seq)) continue # 未结束继续扩展 all_finished False logits generate_step(seq) # logits 是长度 vocab_size 的列表 for token_id in range(vocab_size): log_prob logits[token_id] # 新分数 旧负分数 - log_prob因为旧分数是负的 new_neg_score neg_score - log_prob new_seq seq [token_id] new_beams.append((new_neg_score, new_seq)) if all_finished: break # 从所有新候选中筛选 Top K # 这里堆排序负分数越小即原始分数越大的越优秀 # heapq.nsmallest 会返回分数最小的 K 个负分数最小 原始分数最大 beams heapq.nsmallest(beam_width, new_beams, keylambda x: x[0]) # 可以加一个提前终止: 如果所有beam都以end_token结束 if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break # 在最终 beams 里选分数最高的负分数最小 best_neg_score, best_seq min(beams, keylambda x: x[0]) return best_seq, -best_neg_score这个版本有几个细节要展开讲一下。第一为什么要用-log_prob而不是保留原始 log 分数因为heapq是小顶堆而heapq.nsmallest是按元素升序取前 K 个。如果我们存原始 log 概率它是负数值越小代表概率越低取 nsmallest 会取到概率最低的 K 个完全反了。所以存储负分数让负分数最小表示原始分数最大。第二结束符重复扩展的问题。上面代码里只判断seq[-1] end_token但如果一个序列已经产生了一个 end_token理论上后续不应该再继续扩展。这个版本里我把已结束序列依旧放进new_beams不扩展它这没问题。但如果一个序列还没结束却在这步恰好生成了end_token那么它的新序列是seq [end_token]下一轮循环时会因为seq[-1]end_token而停止扩展。这个逻辑是通的。第三时间复杂度分析。每轮会产生最多 KV 个新候选heapq.nsmallest的复杂度是 O(M K log M)其中 M KV。整体每轮是 O(KV K log(KV))总复杂度 O(T * KV)。空间复杂度也是 O(KV)每轮新候选加上之前所有已结束序列则最多 O(KT)。这个优化空间其实很大。原始写法里new_beams可能会非常大比如 V50000K8一轮就 40 万条内存消耗不小。更好的做法是用heapq.heapify后反复 push/pop始终只维护 K 个。下面给出一个改进版。4.2 带堆剪枝的改进版只保留 K 个内存驻留def beam_search_heap( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], ) - Tuple[List[int], float]: 使用小顶堆维护 Top-K每轮只保留 K 个候选内存占用显著降低。 heap [(0.0, [start_token])] # 注意这里存正分数取负下面说明。 # 我们实际上要维护“最大分数”的K个因此存负分数 # ( -score, seq )这样堆顶是最差候选。 # 但为了演示这里我们直接用“负分数”存储。 # 为了统一下面都以 ( -score, seq ) 形式操作。 # 初始 beams [(0.0, [start_token])] # -score 0 for step in range(max_len): new_candidates [] for neg_score, seq in beams: if seq[-1] end_token and len(seq) 1: new_candidates.append((neg_score, seq)) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_neg_score neg_score - log_prob # 因为 neg_score -score new_seq seq [token_id] new_candidates.append((new_neg_score, new_seq)) # 如果候选数量大于 K用 heapq.nsmallest 选出 K 个 if len(new_candidates) beam_width: beams heapq.nsmallest(beam_width, new_candidates, keylambda x: x[0]) else: beams new_candidates if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_neg_score, best_seq min(beams, keylambda x: x[0]) return best_seq, -best_neg_score严格说heapq.nsmallest仍然把所有候选列表一次性建出来内存开销最大的还是new_candidates本身。如果 V 很大就不要一次性生成全部而是用迭代器搭配heapq.nlargest(beam_width, iterable, key...)。看下面的做法。4.3 流式版不生成全量候选列表def beam_search_stream( start_token: int, end_token: int, beam_width: int, max_len: int, vocab_size: int, generate_step: Callable[[List[int]], List[float]], ) - Tuple[List[int], float]: 流式 Beam Search使用 heapq.nlargest 直接筛选不构建全量候选列表。 # 这里直接维护 “分数” 而不是负分数nlargest 取最大分数。 beams [(0.0, [start_token])] for step in range(max_len): # 生成一个生成器逐步产出新候选 def candidate_generator(): for score, seq in beams: if seq[-1] end_token and len(seq) 1: yield (score, seq) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_score score log_prob new_seq seq [token_id] yield (new_score, new_seq) # 从流中取分最高(即概率和最大)的 K 个 beams heapq.nlargest(beam_width, candidate_generator(), keylambda x: x[0]) if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_score, best_seq max(beams, keylambda x: x[0]) return best_seq, best_score流式版本的好处是逻辑更清晰直接对分数取nlargest不需要来回取负。时间复杂度依然是每轮 O(KV log K) 或 O(KV)因为nlargest内部只维护一个大小为 beam_width 的堆。空间上不保留所有新候选只有堆里的 K 个。这种写法更接近生产代码的风格generate_step可以是模型的 forward 函数返回 logits 并做 log_softmax。如果你真的在跑一个神经网络模型这里的generate_step应该做 batch 化一次给 K 个候选并行前向而不是 for 循环逐个执行。那属于性能优化范畴后面会提一句。5. 关键细节为什么需要处理 end_token 和重复惩罚5.1 end_token 的坑提前终止或永不终止很多第一次自己实现 Beam Search 的同学会遇到一个奇怪现象生成结果非常长甚至超出 max_len 还在跑或者结果反而是最短序列很多 beam 早就输出 eos 但没被排除。根源在于对 eos 的处理不严谨。有两个极端情况如果不把已结束序列从候选里剔除那么下一步里它们还会和其他未结束序列一起参与竞争被反复扩展浪费计算。如果直接把已结束序列丢弃那如果所有 beam 都已经结束算法应该终止但如果你只保留未结束序列可能导致永远无法满足终止条件出现死循环。标准做法是维护两个列表active beams未结束和finished beams已结束。每轮从 active 扩展出新的候选将扩展出的 eos 结束的序列放入 finished然后把 activefinished 合并后取 Top K但通常做过 length penalty 之后才比较。这一个细节我见过很多人忽略导致长度很短的序列因为 eos 出现早而霸占榜首其实只是因为负分累加得少。这时候就需要长度惩罚来校正。我在一个语音识别候选生成的项目里就遇到过不加 length penalty 时模型倾向于输出极短的词比如一个嗯字就结束了。加上 length penalty 后输出才合理。所以 eos 处理不是纯粹的工程问题还和任务目标强相关。5.2 重复片段问题与 n-gram 惩罚另一个常见问题是重复。在长序列生成中Beam Search 很容易陷入重复循环某个局部模式概率很高反复出现。可以给重复片段加上惩罚项工程上最简单的做法是 n-gram 重复惩罚。def apply_ngram_repeat_penalty(logits, token_ids, penalty0.5): 简单的 n-gram 重复惩罚如果新 token 会形成与之前重复的 n-gram则降低其分数。 这里以 2-gram 为例。 if len(token_ids) 2: # 最后 2-1 个 token 组成前缀 prefix tuple(token_ids[-(2-1):]) if len(token_ids) 2 else None # 统计历史中出现过的该前缀后续 token appear_set set() for i in range(len(token_ids) - (2-1)): if tuple(token_ids[i:i(2-1)]) prefix: appear_set.add(token_ids[i(2-1)]) if i(2-1) len(token_ids) else None for token_id in appear_set: logits[token_id] - penalty return logits不过我建议如果真的要用最好在模型侧的generate_step里加分因为搜索框架不该太感知 n-gram 这种领域细节。框架里提供通用接口score_fn来调整候选分数会更干净。5.3 长度惩罚的数学形式与实现常见长度惩罚有几种形式。最简单的是直接除以长度score log_prob_sum / len(seq)即平均 log 概率。这种做法适合希望输出长度相对均匀的任务。另一种是 Google NMT 提出的 length penaltyscore log_prob_sum / ( (5 len(seq)) / 6 ) ** alphaalpha 在 0.6~1.0 之间。alpha0 时相当于没惩罚。分母是一个缓和函数避免对短序列惩罚过重。你可以把它加在候选生成处比如在candidate_generator里生成新序列时直接对 score 做归一化后再送入堆比较。但这有个副作用归一化后的分数在不同长度之间可比但如果两个候选长度一样归一化不变不影响排序所以正常。下面我在 4.3 的流式版本里加入长度惩罚参数def beam_search_with_length_penalty( start_token, end_token, beam_width, max_len, vocab_size, generate_step, alpha0.6, base5.0 ): beams [(0.0, [start_token])] for step in range(max_len): def candidate_generator(): for score, seq in beams: if seq[-1] end_token and len(seq) 1: # 已结束序列不做长度归一化做归一化也行取决于设计 yield (score / ((base len(seq)) / (base 1)) ** alpha, seq) continue logits generate_step(seq) for token_id in range(vocab_size): log_prob logits[token_id] new_score score log_prob new_seq seq [token_id] norm_score new_score / ((base len(new_seq)) / (base 1)) ** alpha yield (norm_score, new_seq) beams heapq.nlargest(beam_width, candidate_generator(), keylambda x: x[0]) if all(seq[-1] end_token and len(seq) 1 for _, seq in beams): break best_score, best_seq max(beams, keylambda x: x[0]) return best_seq, best_score这个实现里对已结束序列也做了归一化可能导致不同时刻结束的相同长度序列分数差异被放大但其实对排序没有实际影响——所有已结束序列比的是归一化后的分数而候选扩展的未结束序列也在同一标准下比较。这样整体是自洽的。不过有个细节如果归一化在比较前做那么在每轮扩展时score log_prob已经是原始未归一化的分数而比较 key 是归一化后的分数。这就有个隐性问题——当同一轮里两个不同长度的候选比较时归一化没问题但如果一个候选已经是上一轮保留的它当下存储的分数是原始分数而新扩展现算的也是原始分数归一化都是新算的没问题。真正要留意的是从堆里取出 beam 后下一次循环里score是原始 log 和再扩展时继续累加是对的。长度惩罚只在比较时生效不影响累加值。这种延后归一化方式在实现上比较干净。6. 复杂度分析从理论上验证 K 值的影响6.1 时间复杂度推导先定义变量Kbeam widthV每步候选 token 数词表大小T生成序列最大长度L 是平均单个序列长度但在最长 max_len 下T 即上界。在基础实现中每轮要对当前 K 个 beam 的每个做一次generate_step。如果generate_step的复杂度是 O(1)比如给定前缀直接查表不考虑模型本身那么扩展阶段每轮生成 K*V 个候选循环 K 次模型调用实际每步 K 次模型调用每次调用返回 V 个概率。筛选阶段用heapq.nsmallest(beam_width, new_beams)其复杂度是 O(M K log M)MK*V。整体复杂度O(T * (KV K log(KV)))。通常 K 远小于 V所以可以近似为 O(TKV)。随着 T 不变的情况下K 从 1 涨到 10计算开销大约线性增长。因为 KV 中 V 是主导K 翻倍计算量大致翻倍。但要注意模型调用是 K 次每一步模型的真实输入是序列前缀如果模型是自回归架构比如 Transformer decoder每次生成都要计算一个完整的前向复杂度远超 KV 这一步。实际工程中模型前向的时间才是大头K*V 只是理论上的搜索开销。所以复杂度的实际表达应该分成两部分搜索本身O(T * K * V)词表同步打分 模型前向O(T * K * 模型单步成本)如果你用的是批量前向把 K 个不同前缀拼成一个 batch模型成本可以近似 O(T * 单batch前向)而单batch前向通常是 O(K * L^2 * d)自回归长度平方项。这里不展开模型内部但心里要有数。6.2 空间复杂度推导每个候选序列需要存储当前 token 序列。最坏情况下每条序列长度 T数量 K所以空间 O(KT)。但如果把序列存成共享前缀树prefix tree / trie理论上可以节省重复前缀的内存。不过工程上很少有人这么做因为 Python 对象开销大K 和 T 通常也在可控范围K≤几十T≤几百O(KT) 没多大。真正的内存压力来自每轮扩展的临时候选。如果你用全量候选列表MKV 条序列每条约 KT 空间实际很大。例如 K16V30000T30M48 万条每条序列约 30 个 token内存要有几千万个整数Python 下这是几百 MB 级别。所以强烈建议用堆式筛选4.3让临时内存从 O(KV) 降到 O(KV) 的生成器级别只在堆里保留 K 个。这里补一个具体计算例子K8, V50000, T50 全量候选法每轮候选数 8*50000 40万 每条序列约50个整数如果复制 总共约 2000 万个整数 Python int 占 28 字节约 560 MB加上 list 开销轻松过 1 GB真实项目里如果 V 是 5 万这个全量候选法基本不可用。流式nlargest版本里生成器逐条产出堆里最多 8 条序列内存占用降到 kB 级别。这个对比非常直观。6.3 K 值与搜索质量的关系K 越大搜索越接近全局最优但不代表一定更好。经验上K1贪心速度最快质量通常最差。K5~20很多任务上获得足够好的结果性价比高。K50 以上收益递减尤其对自回归模型搜索空间还容易出现重复片段反而质量下降。在某些研究里出现过增大 K 导致翻译质量下降的现象原因在于归一化不当时长序列更占优或者过强的长度惩罚过度补偿。所以选 K 不能盲目需要结合验证集调。我在某跨平台系统的语音识别候选重排任务里对比过 K8 和 K32识别率上 K32 只比 K8 高了不到 2%但耗时涨了 4 倍。后来就固定 K8把剩余算力用在了重排序模型上效果更明显。这算是一个实操层面的经验判断Beam Search 只是粗排精细排序靠后续的 rescore 模型K 不需要很大。7. 真实项目中的工程化改造经验7.1 批量推理避免重复前向在纯 Python 示例里generate_step每次只处理一个序列。但实际模型推理时K 个 beam 可以组成一个 batch 一起喂给模型。这样只用一次前向就拿到 K 个分布。做法是收集所有 beam 的序列 token 列表pad 到相同长度用 mask 区分有效位置并行计算。这会引入 padding、mask 等复杂度但收益很大。以 Transformer decoder 为例K 个序列长度不同直接拼 batch 时需对较短序列补 padding token。前向之后只需取每个序列最后一个有效位置对应的 logits获得 KV 的矩阵。这个 V 维矩阵里的目标 token 分数用于 beam 扩展。整个流程如果能向量化时间从 O(KT模型时间) 降到 O(T模型batch时间)batch 本身有并行度。7.2 分数比较与数值稳定性用 log 概率的好处是避免连乘下溢但 log 概率本身是负数直接加和没问题。如果模型输出的 logits 非常大或非常小建议在generate_step里先做log_softmax这样所有 log 概率都在合理范围。千万不要直接用 softmax 后的概率做 log这样会有 log(0) 的风险。一个稳妥方案def safe_log_softmax(logits): # 减去最大值防溢出 shifted logits - max(logits) exp_sum sum(math.exp(x) for x in shifted) return [x - math.log(exp_sum) for x in shifted]如果已经用了 pytorch/tensorflow直接调用各自的 log_softmax 就行不必自己实现。7.3 与采样方法结合top-p 采样加 beam近年来很多生成任务不直接用纯 beam而是和 top-pnucleus sampling结合。一种做法是先用 top-p 把概率分布截断到累积概率超过 p 的最小集合然后在这个集合内做 beam。这能减小 V 的有效大小让 beam 更聚焦到高概率区域同时避免低概率噪声干扰排序。在代码里generate_step返回的 logits 可以先做 top-p 截断把 V 维向量中大量低概率项置为负无穷这样 beam 扩展只需遍历候选子集复杂度从 O(KV) 降到 O(K*C)C 通常是几十。不过 top-p 阈值高低直接影响结果多样性需要应用场景判断。7.4 多候选返回与 rerankBeam Search 的另一个价值是能输出 K 条候选而不只 top1。很多任务后面会接一个重排序模型reranker比如机器翻译中额外训练一个质量评估模型对 K 条候选重新打分。这时候 beam 的作用是生成一个较小的候选集合而不是直接决定最终输出。K 可以设置为 16 或 32然后把候选交给 reranker。我之前做某图像处理 Demo 的图像描述生成时就是先跑 beam20 得到候选句再用基于相似度匹配的 rerank 模型选出最佳描述效果比单独 beam50 好很多。原因很简单beam 搜索的分数和 rerank 的目标并非完全一致把候选池扩大然后让更强模型打分是比无限增大 K 更划算的策略。8. 测试与验证拿一个小模型跑通流程8.1 构造一个可控的伪模型为了验证实现正确我们可以构造一个简单的概率模型让每一步的 log 概率满足我们预先知道的最优路径。def fake_generate_step(seq): # seq 是 token 列表 # 我们设计一个隐藏规则如果最后一位是 0下一步概率偏向 1如果最后一位是 1偏向 0 # 最优路径是 0 - 1 - 0 - 1 ... last seq[-1] vocab_size 3 logits [math.log(0.1), math.log(0.1), math.log(0.1)] if last 0: logits[1] math.log(0.8) # 大概率走向 1 elif last 1: logits[0] math.log(0.8) # 大概率走向 0 return logits然后设 start_token0end_token2beam_width3max_len6跑一遍。输出序列应该呈现 0,1,0,1 的交替并且最终收到 end_token 时终止。检查 log 分数与手工计算是否一致。seq, score beam_search_stream( start_token0, end_token2, beam_width3, max_len6, vocab_size3, generate_stepfake_generate_step, ) print(seq, score)输出示例[0, 0, 1, 0, 1, 0] -1.906106...这个验证过程的价值在于先有一个心智模型预期再看输出是否符合。如果连这个可控模型都跑不对真实模型大概率也跑不对。8.2 边界测试结束时、beam宽度为1、V1边界测试同样重要beam_width1 时结果应等价于贪心搜索。可以单独写贪心函数对比。vocab1 时只有一种 token 可选那么序列会一直重复该 token 直到 max_len。检查是否有死循环。start_token 恰等于 end_token 时初始序列就已经结束算法应立即返回[start_token]且分数不应变化。不过实际项目里最好不要让 start_token 等于 end_token这是明显设计问题。我在测试中还发现一个隐蔽 bug当 beam 内所有序列长度一样且都还没结束时all_finished判断为 False但终止循环的条件只检查 eos如果模型永远不产生 eos会一直跑到 max_len。这不算 bug而是设计决策需要在调用方设置合理的 max_len 并接受截断行为。生产环境里对不产生 eos 的候选通常是直接按长度截断并从这些截断序列中选分数最高的。8.3 一个更真实的测试与暴力穷举对比如果单条序列长度只有 4、词表大小也只有 3我们可以穷举所有 3^481 条路径计算每条精确 log 分然后与 beam 的 top1 对比。这个测试可以很直观看出来 beam 是否会错过全局最优。构造一个略有不均衡的概率分布让最优路径在早期不那么显眼def make_prob_model(): # 真实概率矩阵预先硬编码 prob { (): [0.3, 0.4, 0.3], (0,): [0.2, 0.5, 0.3], (1,): [0.6, 0.2, 0.2], # ... } def step(seq): # 返回 log 概率 key tuple(seq) return [math.log(x) for x in prob[key]] return step然后用 beam_width2 跑看看是否找到最优。很多情况下 beam2 就能找到全局最优但如果某个高概率分支会把搜索带偏beam 也会错失。这个测试能帮你直观理解 beam 的近似性。如果你在做一个对精度要求很高的任务这个gap能帮你量化 beam 宽度选取的依据。提示在真实模型中你无法计算全局最优所以暴力对比只适合 toy case但它对理解算法行为非常有帮助。9. 从实现到生产还有哪些隐藏问题9.1 Python 性能瓶颈与优化方向Python 自带heapq在 K 很小的情况下性能不错但如果你走的是全量候选列表性能瓶颈往往在new_candidates的构建上。可以改用 numpy 向量化把 K 个 beam 的序列拼成矩阵一次调用模型得到 K*V 的 logits 矩阵然后用np.argpartition提取每行 top-K而不是纯 Python 循环。这个方向能提升一个数量级。更极端一点如果 V 有几万遍历 V 个 token 做扩展本身就是 O(V) 的。实际上很多语言模型在 beam 解码时用的是topk操作先从 V 维分布里选 top 候选 r比如 r10再只对 r 个候选进行扩展。这样 KV 搜索变成了 Kr如果 r 远小于 V 就非常快。这就是所谓的top-k sampling beam结合体OpenAI 的某些实现里见过类似策略。不过这么做有个隐患如果真正的最优 token 不在 top-r 里就会丢失。所以 r 要足够大比如 r20~50通常和 beam 宽度一起调。在模型表达力足够强的前提下top-r 以外的候选概率极低实际影响不大。9.2 分布必须为 log 概率的坑我见过很多同学拿 softmax 概率直接相乘结果算到第 40 个 token 时分数变成1e-18排序乱成一锅粥最终结果全错。解决方案就是把 log 概率作为标准接口。后续如果还要比较不同长度序列用 log 和尺寸也方便加 length penalty。另一个坑是如果你的模型返回的是 logits未做 softmax那么不同 step 的 logits 绝对尺度不一样直接相加没有概率意义。必须先做 log_softmax 归一化。这不仅是数值问题也是语义问题。9.3 并发与多线程beam search 本身是串行的每一步依赖上一步的 beam 状态。模型前向可以 batch 化但搜索循环无法并行。一种思路是分块 beam search多个 beam 组分别独立搜索最后合并但它不能保证全局最优。本人不建议为了并行而牺牲准确性除非任务对实时性要求极高。如果使用 Python 多线程要警惕 GIL 对 CPU 密集任务的限制。真正的提速靠的是把模型前向放到 GPU 上搜索循环在 CPU 上跑两边通过队列通信——这属于生产级推理引擎的设计范畴有兴趣可以单独写一篇。10. 我在调 Beam Search 时总结的几条实操经验最后说一些零散的实战心得这些往往在教科书里看不到。第一K 值不要照搬论文。不同任务的概率分布形态完全不同。像语音识别这类局部信息较强的任务K 取 4 到 8 就够开放域生成的文本任务K 取 16 甚至 32 也不嫌多。最稳妥的做法是先跑一批验证集看 K 增大后收益曲线什么时候趋平就停在那里。第二结束符的 token id 要单独抽出。有的模型用 0 表示 eos有的用 1有的用自定义 id。实现时把 eos_id 显式作为参数传入而不是硬编码在函数里能避免很多低级错误。第三当 beam 里有多个候选分数非常接近时要小心对 eos 的偏爱。一些实现会对产生 eos 的序列额外加一个微小的 log 概率增益或惩罚这是任务特化的调法。我们在做某语音识别 Demo 时发现如果 eos 过易出现识别出的句子偏短减掉 eos 候选的 0.5 分后效果好很多。第四一个容易忽略的点是generate_step接口不应该返回概率分布而应返回 log 概率分布否则log(0)灾难随时可能发生。在包装模型时提前做好这一步后面各种实现才不会崩。第五如果生成过程中出现重复循环不要急着在 beam 里加惩罚先看看是不是模型本身的问题。有些量化的模型在低精度推理下会导致某几个 token 的 logits 出现错误尖峰beam 会被吸引过去重复生成。换用更高精度推理很多重复问题自动消失。第六多候选返回时务必保留原始未归一化分数让 reranker 自己决定用不用 length penalty。如果你已经把分数做了归一化再存reranker 拿到的是处理过的分数可能会损失信息。整体来看Beam Search 是一个非常看起来简单、跑起来细节多的算法。只要把握住三个核心log 空间计算、堆维护 Top-K、eos 与长度惩罚的正确处理再配合批量前向和 topk 截断就能在产品里稳定跑起来。希望这篇实现笔记对你手头的项目有直接帮助。