嗨，各位技术小伙伴们！今天咱们来聊一个在自然语言处理（NLP）和序列生成任务中超重要的算法——集束搜索（Beam Search）！🎯 无论是机器翻译、文本摘要，还是对话系统，集束搜索都能让AI生成的句子更通顺、更合理。它到底是怎么工作的？和贪心搜索（Greedy Search）有什么区别？别急，咱们就通过下面文章一次性搞懂！📚
🌰 开篇小例子：翻译“Hello”到中文
假设我们训练了一个中译英的神经网络模型，输入是“我爱你”，模型会逐步生成英文词：
- 第一步：生成“I”、“LOVE”、“YOU”（概率分别为0.3、0.6、0.1）。
- 第二步：根据上一步选的词，继续生成下一个词……
如果用贪心搜索，在翻译每个字的时候，直接选择条件概率最大的候选值作为当前最优。
而集束搜索是对贪心算法的一个改进算法。相对贪心算法扩大了搜索空间🎉
🤖 什么是集束搜索？
集束搜索是一种启发式搜索算法，它在每一步生成序列时，保留概率最高的前k个候选序列（k称为“集束宽度”），然后继续扩展这些序列，直到生成完整结果。
🔄 和贪心搜索的区别
| 方法 | 每一步候选数 | 优点 | 缺点 |
|---|---|---|---|
| 贪心搜索 | 1 | 速度快 | 容易陷入局部最优（如“你你好”） |
| 集束搜索 | k（可调） | 能找到全局更优的序列 | 计算量比贪心搜索大 |
💻 集束搜索的步骤（以机器翻译为例）
假设我们要将英文“I love NLP”翻译成中文，集束宽度 k=2：
📌 步骤1：初始化
- 输入：
<BOS>（句子开始标记）。 - 当前候选序列：
[<Bos>]（概率=1.0）。
📌 步骤2：扩展序列
- 第一步扩展：
- 模型预测下一个词的概率：
"我": 0.7,"你": 0.2,"他": 0.1。 - 保留前k=2个候选：
- 序列1：
[<Bos>, "我"]，概率=1.0 * 0.7 = 0.7 - 序列2：
[<Bos>, "你"]，概率=1.0 * 0.2 = 0.2
- 序列1：
- 模型预测下一个词的概率：
- 第二步扩展：
- 对序列1扩展：
- 预测下一个词：
"喜欢": 0.6,"爱": 0.3,"讨厌": 0.1。 - 新候选：
[<Bos>, "我", "喜欢"]，概率=0.7 * 0.6 = 0.42[<Bos>, "我", "爱"]，概率=0.7 * 0.3 = 0.21
- 预测下一个词：
- 对序列2扩展：
- 预测下一个词：
"喜欢": 0.4,"爱": 0.5,"讨厌": 0.1。 - 新候选：
[<Bos>, "你", "爱"]，概率=0.2 * 0.5 = 0.1[<Bos>, "你", "喜欢"]，概率=0.2 * 0.4 = 0.08
- 预测下一个词：
- 合并所有候选，保留前k=2个：
[<Bos>, "我", "喜欢"]（0.42）[<Bos>, "我", "爱"]（0.21）
- 对序列1扩展：
- 第三步扩展（直到遇到
<EOS>结束标记）：- 假设最终选概率最高的序列：
[<Bos>, "我", "爱", "NLP", <EOS>]→ “我爱NLP”。
- 假设最终选概率最高的序列：
🚀 应用示例：用PyTorch实现集束搜索
以下是一个简化的机器翻译集束搜索代码示例（假设模型已训练好）：
import torch
import torch.nn as nn
# 模拟翻译模型（实际中替换为你的模型）
class DummyTranslator(nn.Module):
def __init__(self):
super().__init__()
self.vocab_size = 1000 # 假设词表大小为1000
self.max_len = 20 # 最大生成长度
def forward(self, input_ids, past_key_values=None):
# 模拟输出：每一步生成logits（未归一化的概率）
batch_size = input_ids.shape[0]
logits = torch.randn(batch_size, self.vocab_size) * 0.1 # 随机生成logits
return logits, None
# 集束搜索函数
def beam_search(model, input_ids, beam_width=3, max_len=20):
# 初始化：当前候选序列和它们的概率
sequences = [[input_ids[0].tolist()]] # 初始序列（假设batch_size=1）
scores = [0.0] # 初始概率（对数概率）
for _ in range(max_len):
all_candidates = []
for seq, score in zip(sequences, scores):
# 如果序列已结束（遇到<EOS>），跳过扩展
if seq[-1] == 2: # 假设2是<EOS>的ID
all_candidates.append((seq, score))
continue
# 用模型预测下一个词的概率
input_tensor = torch.tensor([seq[-1]]).unsqueeze(0) # 模拟输入（实际需更复杂）
logits, _ = model(input_tensor)
probs = torch.softmax(logits, dim=-1)[0].tolist() # 转换为概率
# 生成所有可能的候选序列
for i in range(len(probs)):
new_seq = seq + [i]
new_score = score + torch.log(torch.tensor(probs[i])).item() # 累加对数概率
all_candidates.append((new_seq, new_score))
# 按概率排序，保留前beam_width个候选
ordered = sorted(all_candidates, key=lambda x: x[1], reverse=True)
sequences = [seq for seq, score in ordered[:beam_width]]
scores = [score for seq, score in ordered[:beam_width]]
# 如果所有序列都结束了，提前终止
if all(seq[-1] == 2 for seq in sequences):
break
# 返回概率最高的序列
best_seq = sequences[0]
return best_seq
# 测试
model = DummyTranslator()
input_ids = torch.tensor([[1]]) # 假设1是<BOS>的ID
output_seq = beam_search(model, input_ids, beam_width=3)
print("Generated sequence:", output_seq) # 输出类似 [1, 10, 20, 2]（<BOS>, 词1, 词2, <EOS>）代码说明：
DummyTranslator是一个模拟的翻译模型，实际使用时替换为你的PyTorch/TensorFlow模型。beam_search函数实现了集束搜索的核心逻辑：- 每一步扩展所有候选序列。
- 用对数概率累加避免数值下溢。
- 保留概率最高的前k个序列。
- 最终返回概率最高的完整序列。
📊 集束搜索的优缺点
| 优点 | 缺点 |
|---|---|
| 能找到全局更优的序列 | 计算量比贪心搜索大 |
| 适合序列生成任务（如翻译、对话） | 需要调参（集束宽度k） |
| 可结合长度惩罚（避免生成过短序列） | 可能仍陷入局部最优（k较小时） |
长度惩罚（Length Penalty）：
为了平衡序列长度和概率，可以引入长度惩罚项：
其中 α 是超参数（通常0.6~1.0），避免模型倾向于生成短序列。
💡 总结与建议
- 什么时候用集束搜索？
- 需要生成合理序列的任务（翻译、摘要、对话、文本生成）。
- 贪心搜索效果不佳时（如生成重复或不通顺的句子）。
- 如何选择集束宽度k？
- 小k（如2~5）：速度快，适合实时应用。
- 大k（如10~20）：质量更高，但计算量大。
- 实际中可通过验证集调参。
- 进阶优化：
- 结合Top-k采样或核采样（Nucleus Sampling）增加多样性。
- 使用Transformer+Beam Search（如Hugging Face的
generate方法）。
希望这篇博客能帮你彻底理解集束搜索！如果有任何问题或想看的应用场景，欢迎在评论区留言~记得关注哦！ 🌟