题目重点：以透明的固定宽度训练契约实现 beam search，解释完成序列、平局、长度偏置和无法保证全局最优。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 训练版 `next_tokens(prefix)` 返回 token 与其条件 log probability；每轮按原始累计 log probability 只保留 beam_width 条路径，平局按 token 序列字典序；EOS 代表结束且不返回在文本中。
- 最多扩张 max_steps 轮，结束路径仍与未结束路径竞争同一 beam；不加长度归一化或多样性惩罚。真实解码器需要与面试官确认这些规则。
解题路线
- 初始 beam 是空前缀，得分 0。每轮只扩张未结束路径，将下一个 token 的 log probability 加到累计分数；结束路径原样保留。
- 把候选按得分降序和稳定平局规则排序，截取前 B 条。若全部结束可提前停止，未结束且达到最大步数则返回当前最佳前缀。
- 对数域避免逐 token 概率连乘下溢；这里的原始 logprob 会偏向短序列，若要 length penalty 应明确其公式以及与 EOS 比较方式。
- beam search 是启发式宽度剪枝，不保证全局最优：早期被剪掉的低概率前缀可能在后续超过保留前缀。用独立逐层展开参考实现验证 800 组宽度/步数/分数设置。
手动走一遍
B=2，首步 a 得 −0.1、b 得 −0.3；次步各自扩张，再从所有未结束扩张和已结束路径中选得分最高两条。若某前缀得到 EOS，之后不再向它添加 token，但仍参与排名。数字只示意算法而非原题给定概率。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
def beam_search(next_tokens, beam_width: int, max_steps: int, eos: str = '<eos>'):
"""next_tokens(prefix) -> iterable (token, log_probability).
Keep at most beam_width hypotheses after each step by raw cumulative log
probability. Ties use lexical token sequence. Return best (tokens, score).
"""
if beam_width < 1 or max_steps < 0:
raise ValueError('invalid search budget')
beams = [((), 0.0, False)]
for _ in range(max_steps):
candidates = []
for prefix, score, finished in beams:
if finished:
candidates.append((prefix, score, True))
continue
for token, log_probability in next_tokens(prefix):
if not isinstance(token, str) or token == '':
raise ValueError('tokens must be nonempty strings')
candidates.append((prefix + (token,), score + log_probability,
token == eos))
if not candidates:
break
candidates.sort(key=lambda item: (-item[1], item[0]))
beams = candidates[:beam_width]
if all(finished for _, _, finished in beams):
break
best = min(beams, key=lambda item: (-item[1], item[0]))
tokens = best[0][:-1] if best[2] and best[0][-1] == eos else best[0]
return tokens, best[1]
复杂度
每步每条 beam 最多扩张 V 个 token，T 步排序/取前 B 约 O(T·BV log(BV)) 时间，保存候选 O(BV·T)（若复制整个前缀）；实际模型计算成本另计，可用堆和共享前缀降低开销。
易错点与追问
容易失分的边界
- 用普通概率连续相乘导致长序列下溢。
- 已经 EOS 的路径继续扩张。
- 每一轮只从上一轮最优单路径扩张，退化成 greedy decoding。
- 声称 beam search 一定找到全局最可能完整序列。
- 未问 length penalty 就把短序列偏置当成模型偏好。
继续深挖
- beam_width=1 与 greedy decoding 何时等价？
- 如何加入长度惩罚而不错误比较已结束和未结束路径？
- top-k/top-p sampling 与 beam search 的目标有什么差异？
- 模型词表很大时如何减少扩张/排序成本？
