-
算法名称和简介
带长度惩罚的Beam Search解码（批量版）
在自回归生成任务（机器翻译、文本生成、语音识别等）中，使用Beam Search在每步扩展多个候选，结合长度惩罚α缓解短句偏好，支持重复n元阻断（no-repeat n-gram）避免循环，结合温度缩放与top-k过滤控制多样性与计算开销。算法支持批量解码与EOS早停，返回每个样本的前N个候选序列及其对数概率。
-
输入输出说明
输入
-
ModelStep: 增量解码接口，签名：
logits, new_cache = ModelStep(prev_tokens, cache, batch_beam_mask)
- prev_tokens: [B_active]，当前每个活跃beam的最后一个token（或用于增量的当步输入）
- cache: 模型缓存（见“关键状态缓存”）
- batch_beam_mask: [B_active]，可选，用于指示哪些beam是有效（便于一次性批处理）
- 返回 logits: [B_active, V]（未归一化logits），new_cache 为更新后的缓存
-
prompts: List[List[int]]，长度为batch_size，每项为初始提示序列（可含BOS）
-
max_steps: 最大解码步数T_max（不含prompt长度）
-
beam_size: 每个样本的beam数B
-
n_best: 返回的候选数N（N ≤ B）
-
alpha: 长度惩罚系数α（GNMT样式）
-
no_repeat_ngram_size: n元阻断大小，n=0表示关闭
-
temperature: 温度τ > 0，τ=1为不缩放
-
top_k: k值（1 ≤ k ≤ V），k=V表示不裁剪
-
eos_token_id, bos_token_id, pad_token_id: 特殊符号ID
-
early_stop: 是否启用EOS早停（默认true）
输出
函数定义
function LengthPenalty(L, alpha):
// GNMT长度惩罚
// lp(L) = ((5 + L) ^ alpha) / (6 ^ alpha)
return ((5 + L) ^ alpha) / (6 ^ alpha)
function TopKFilter(logits_vec, k):
// logits_vec: [V]
if k >= V: return logits_vec
indices = argsort_desc(logits_vec) // 从大到小
threshold_index = indices[k-1]
threshold = logits_vec[threshold_index]
for v in 0..V-1:
if logits_vec[v] < threshold:
logits_vec[v] = -INF
return logits_vec
function NoRepeatNGramMask(seq, n, V):
// 返回 banned_mask: [V]，true表示禁止
if n <= 0 or length(seq) < n - 1:
return BoolVector(V, false)
prefix = seq[(length(seq) - (n - 1)) : length(seq)] // 长度n-1
banned = Set() // 将与该prefix构成重复n-gram的下一个token
// 扫描历史，收集所有prefix后续token
for i in 0 .. (length(seq) - n):
if seq[i : i + (n - 1)] == prefix:
next_token = seq[i + (n - 1)]
banned.add(next_token)
banned_mask = BoolVector(V, false)
for t in banned:
banned_mask[t] = true
return banned_mask
数据结构
struct BeamState:
tokens: List[int] // 已生成（不含prompt）；用于输出可与prompt拼接
last_token: int
sum_logprob: float // 累积对数概率（含当前tokens）
length: int // 已生成长度（含EOS时计入）
cache: Any // 模型增量缓存（KV缓存/隐藏状态等）
finished: bool // 是否已生成EOS
score: float // 用于排序：sum_logprob / LengthPenalty(length, alpha)
主过程
function BeamSearchBatch(ModelStep, prompts, max_steps, beam_size, n_best,
alpha, no_repeat_ngram_size, temperature, top_k,
eos_token_id, bos_token_id, pad_token_id, early_stop):
batch_size = length(prompts)
V = vocabulary_size()
// 初始化每个样本的beam
beams = List[batch_size] of List[BeamState]
finished_lists = List[batch_size] of PriorityQueue(desc_by_score) // 存储完成的候选
for b in 0 .. batch_size-1:
// 初始化cache：将prompt编码至cache。具体取决于ModelStep/模型结构
init_cache = InitCacheFromPrompt(prompts[b])
init_last = prompts[b][-1] // prompt的最后一个token作为第一步输入
beams[b] = []
// 将单一初始beam复制为beam_size份（共享或拷贝cache）
for k in 0 .. beam_size-1:
beams[b].append(BeamState(tokens=[], last_token=init_last,
sum_logprob=0.0, length=0, cache=Clone(init_cache),
finished=false, score=0.0))
step = 0
while step < max_steps:
step += 1
// 汇总所有未完成beam用于一次性前向
active_indices = [] // (b, j)对
inputs = [] // 对应last_token
caches = [] // 对应cache
for b in 0 .. batch_size-1:
for j in 0 .. beam_size-1:
if not beams[b][j].finished:
active_indices.append((b, j))
inputs.append(beams[b][j].last_token)
caches.append(beams[b][j].cache)
if length(active_indices) == 0:
break // 所有样本都已完成
// 一次性模型前向
logits_batch, new_caches = ModelStep(inputs, caches, None) // 维度: [B_active, V]
// 对每个活跃beam进行温度缩放、top-k、n元阻断，然后log-softmax
log_probs_batch = Matrix(len(active_indices), V)
for idx in 0 .. len(active_indices)-1:
b, j = active_indices[idx]
logits = logits_batch[idx]
// 温度缩放
if temperature > 0:
logits = logits / temperature
// top-k过滤
logits = TopKFilter(logits, top_k)
// n元阻断
if no_repeat_ngram_size > 0:
// 计算基于prompt+当前beam.tokens的序列
prefix_full = Concat(prompts[b], beams[b][j].tokens)
banned_mask = NoRepeatNGramMask(prefix_full, no_repeat_ngram_size, V)
for v in 0 .. V-1:
if banned_mask[v] == true:
logits[v] = -INF
// 归一化为log概率
log_probs = LogSoftmax(logits) // 数值稳定
log_probs_batch[idx] = log_probs
// 为每个样本收集候选扩展，并选出下一个step的beam
new_beams = List[batch_size] of empty list
for b in 0 .. batch_size-1:
// 汇总该样本的所有候选 (来自其未完成beam)
candidate_scores = [] // (score, from_beam_j, token_id, new_sum_logprob, new_cache_idx)
// 记录活跃beam在active_indices中的位置
pos_list = []
for idx in 0 .. len(active_indices)-1:
bb, jj = active_indices[idx]
if bb == b:
pos_list.append((idx, jj))
// 为该样本每个活跃beam枚举所有词汇的扩展（可只枚举非-INF部分）
for (idx, jj) in pos_list:
if beams[b][jj].finished:
continue
log_probs = log_probs_batch[idx] // [V]
for v in 0 .. V-1:
if log_probs[v] == -INF:
continue
new_sum = beams[b][jj].sum_logprob + log_probs[v]
new_len = beams[b][jj].length + 1
lp = LengthPenalty(new_len, alpha)
new_score = new_sum / lp
candidate_scores.append((new_score, jj, v, new_sum, idx))
// 将EOS扩展放入完成列表，不进入下一步beam；非EOS扩展用于下一步beam选取
// 拆分候选
eos_candidates = []
cont_candidates = []
for c in candidate_scores:
(sc, from_j, tok, new_sum, idx) = c
if tok == eos_token_id:
eos_candidates.append(c)
else:
cont_candidates.append(c)
// 将EOS候选加入完成队列（按score排序）
for c in eos_candidates:
(sc, from_j, tok, new_sum, idx) = c
// 生成完整tokens：旧tokens + [EOS]
seq_tokens = Copy(beams[b][from_j].tokens)
seq_tokens.append(tok)
hyp = BeamState(tokens=seq_tokens,
last_token=tok,
sum_logprob=new_sum,
length=beams[b][from_j].length + 1,
cache=null, finished=true, score=sc)
finished_lists[b].push(hyp)
// 从非EOS候选中选出top beam_size个用于下一步
// 注意：若非EOS候选少于beam_size，用EOS结束的候选填充将导致后续无活跃beam，依赖外层循环终止
cont_candidates = top_k_tuples_by_first_desc(cont_candidates, beam_size)
// 组装new_beams[b]
for rank in 0 .. length(cont_candidates)-1:
(sc, from_j, tok, new_sum, idx) = cont_candidates[rank]
old = beams[b][from_j]
new_beam = BeamState(tokens=Copy(old.tokens),
last_token=tok,
sum_logprob=new_sum,
length=old.length + 1,
cache=new_caches[idx],
finished=false,
score=sc)
new_beam.tokens.append(tok)
new_beams[b].append(new_beam)
// 若不足beam_size，补齐占位（将其标为finished以加速结束）
while length(new_beams[b]) < beam_size:
// 尝试从当前最佳EOS候选复制一个“冻结beam”
if length(eos_candidates) > 0:
best_eos = argmax_by_first(eos_candidates)
(sc, from_j, tok, new_sum, idx) = best_eos
old = beams[b][from_j]
frozen = BeamState(tokens=Copy(old.tokens), last_token=tok,
sum_logprob=new_sum, length=old.length + 1,
cache=null, finished=true, score=sc)
frozen.tokens.append(tok)
new_beams[b].append(frozen)
else:
// 没有可用候选时，放入一个PAD占位（不会再被使用）
pad_beam = BeamState(tokens=[], last_token=pad_token_id,
sum_logprob=-INF, length=0, cache=null,
finished=true, score=-INF)
new_beams[b].append(pad_beam)
// 提前终止判断（每个样本独立）
// 若已收集到n_best个完成候选，并且继续扩展的最佳上界不可能超过已完成的最差分数，则早停
all_done = true
for b in 0 .. batch_size-1:
beams[b] = new_beams[b] // 更新beam
// 判断是否还有未完成beam
sample_has_active = false
best_active_upper = -INF
for j in 0 .. beam_size-1:
if not beams[b][j].finished:
sample_has_active = true
// 上界：当前score即为上界（未来对数增量≤0，且长度惩罚非减）
upper = beams[b][j].sum_logprob / LengthPenalty(beams[b][j].length, alpha == 0 ? 0 : alpha)
if upper > best_active_upper:
best_active_upper = upper
if not sample_has_active:
// 无活跃beam，认为此样本已完成
// 若还没有n_best个完成候选，则后续在收尾阶段补齐
continue
all_done = false
if early_stop and finished_lists[b].size() >= n_best:
// 已完成候选中的第N名分数
worst_needed = kth_score_desc(finished_lists[b], n_best)
if best_active_upper <= worst_needed:
// 将所有活跃beam丢弃，标记此样本完成
for j in 0 .. beam_size-1:
beams[b][j].finished = true
if all_done:
break
// 收尾：若完成候选不足n_best，用当前活跃beam按score补足
results = List[batch_size] of empty list
for b in 0 .. batch_size-1:
temp_list = []
// 先取完成候选
while finished_lists[b].size() > 0:
temp_list.append(finished_lists[b].pop()) // 已按score降序
// 不足则从活跃beam补足（按当前score排序）
if length(temp_list) < n_best:
active_sorted = sort_desc_by(beams[b], key=beam.score)
for t in 0 .. min(n_best - length(temp_list), length(active_sorted)) - 1:
temp_list.append(active_sorted[t])
// 只保留前n_best
temp_list = temp_list[0 : min(n_best, length(temp_list))]
// tokens前可选择是否拼接prompt（依任务需求）
results[b] = temp_list
return results
-
关键参数说明
-
beam_size B: 每步保留的候选数，B越大搜索越广但计算更大。
-
n_best N: 返回的候选数，通常N ≤ B。
-
alpha: 长度惩罚强度，α=0时无惩罚；常用0.6左右。采用GNMT形式 lp(L)=((5+L)^α)/(6^α)。
-
no_repeat_ngram_size n: n元重复阻断大小；n=2为常用（阻断重复bigram）。
-
temperature τ: 温度缩放，τ<1使分布更尖锐，τ>1更平坦。
-
top_k k: 每步仅保留概率最高的k个token参与归一化与扩展，k减小可降计算并增多样性；k=V表示不裁剪。
-
early_stop: 启用EOS早停；当已收集到N个完成候选且活跃beam的最优上界不可能超过已完成的第N名分数时，提前终止该样本的搜索。
-
算法复杂度分析
-
时间复杂度（单样本）: O(T · B · V)
- 每步对B个beam计算V维logits并做log-softmax/top-k与阻断；若top-k生效且实现为截断，则有效复杂度近似为O(T · B · k)。
-
批量下总体: O(batch_size · T · B · V)。
-
空间复杂度（单样本）:
- 序列缓存：O(B · T)（存tokens与n-gram索引）
- 模型状态缓存：O(B · CacheSize)
- Transformer解码器KV缓存：O(Layers · Heads · B · T · d_head)
- RNN隐藏态：O(Layers · B · d_hidden)
- 临时logits/log_probs：O(B · V)
-
关键状态缓存
- 编码器输出（encoder memory）：每样本一次性缓存，供所有beam共享。
- 解码器增量缓存（KV缓存/隐藏态）：每个beam一份，随步数累积；扩展到新beam时复制或引用增量更新的缓存。
- Beam层面缓存：
- sum_logprob、length、last_token、finished标志
- tokens（或环形缓冲用于尾部n-gram计算）
- n元阻断索引（可按prefix映射到被禁token集合，或按需即时构建）
-
适用场景说明
-
机器翻译、文本生成、摘要、对话、语音识别等自回归序列生成。
-
需要在质量与速度之间折中（通过B、k、τ、α等调参），并避免重复模式（n元阻断）。
-
批量生产环境推理与服务端在线推理（依赖增量缓存高效迭代）。
-
自检用例（可复现实验）
目的：验证批量、长度惩罚、n元阻断、温度与top-k、EOS早停与N-best排序是否生效。
构造玩具模型
- 词表V=5: {BOS=0, A=1, B=2, C=3, EOS=4}
- 模型规则（确定性logits，方便计算）：
- 若上一个token是BOS(0): logits = [ -10, 3.0, 2.0, 1.0, 0.5 ] // A最优，其次B、C、EOS
- 若上一个token是A(1): logits = [ -10, 0.5, 1.0, 0.0, 2.5 ]
- 若上一个token是B(2): logits = [ -10, 1.5, 0.5, 0.0, 1.0 ]
- 若上一个token是C(3): logits = [ -10, 0.0, 1.0, 0.5, 2.0 ]
- 若上一个token是EOS(4): logits = [ -INF, -INF, -INF, -INF, -INF ] // 不再扩展
- ModelStep实现：读取inputs[i]→按上表返回logits[i], new_cache原样返回。
测试1：基本功能与长度惩罚
- 参数：batch_size=1, prompts=[[BOS]], B=3, N=2, α=0.6, n=0, τ=1.0, k=5, max_steps=4, early_stop=true
- 期望：
- 第一步优先扩展A、B、C
- 后续倾向尽早生成EOS，但长度惩罚抑制过短句偏好（α>0）
- 返回前2个候选，含tokens、sum_logprob、score，且score为降序
测试2：top-k与温度
- 参数：同测试1，但τ=0.7, k=2
- 期望：每步仅保留top-2（例如从BOS只保留A、B），生成多样性受τ影响；结果不含被裁掉的token扩展。
测试3：no-repeat bigram
- 参数：batch_size=1, prompts=[[BOS, A]], B=3, N=2, α=0.0, n=2, τ=1, k=5, max_steps=5
- 期望：若历史出现过某bigram，则阻断再次出现同bigram的后续token；验证NoRepeatNGramMask对当前prefix生效。
测试4：批量与EOS早停
- 参数：batch_size=2, prompts=[[BOS], [BOS, A]], B=2, N=1, α=0.6, n=0, τ=1, k=5, max_steps=5, early_stop=true
- 期望：两个样本独立早停；当已获得N个完成候选且活跃beam上界不超过已完成最差分时，提前结束该样本。
可选断言
- 对于任意候选：score = sum_logprob / LengthPenalty(length, α)
- 对于每步：每样本的活跃beam数≤B
- 被n元阻断的token其log_probs应为-INF
- 若early_stop触发：best_active_upper ≤ 第N个完成候选的分数
备注与实现提示
- 为避免对所有V做log_softmax，可在top-k后仅对k个保留项归一化，并将其余置为-INF。
- 为加速n元阻断，可维护哈希：key为(n-1)-gram前缀，value为其后续token集合；每步仅针对当前前缀查询而非重扫历史。
- 对缓存复制（cache clone）可采用“写时复制”（copy-on-write）或结构化共享，避免大规模内存拷贝。
- 若需要返回不含EOS的序列，可在收尾阶段去除尾部EOS但保留其贡献的对数概率与长度惩罚计算。
