广告位 · 页头(后台-模块管理-新建模块后替换此区域)

推理模型为什么需要搜索策略?Beam Search与蒙特卡洛树搜索如何选择落地

推理模型在回答问题或生成代码时,通常不是一次性输出全部内容,而是逐token或逐步生成。每一步模型都会给出一个概率分布,如果只选择概率最高的那个token,就是贪心解码。贪心解码速度很快,但在数学证明、复杂规划等需要多步推理的任务中,前一步看似最优的选择可能把后续推导带入死胡同。等模型发现错误时,已经无法回退,只能继续生成一个错误答案。这个问题促使研究者把传统搜索算法引入推理过程,Beam Search和蒙特卡洛树搜索就是两种代表性方案。

推理模型为什么需要搜索策略?Beam Search与蒙特卡洛树搜索如何选择落地

一、Beam Search:用固定宽度保留多条候选路径

Beam Search可以理解为贪心解码的扩展。贪心解码每步只保留一个得分最高的序列,而Beam Search每步保留固定数量的候选序列,这个数量称为束宽。算法开始时只有一个输入序列,每步对当前所有候选序列分别计算下一步token的概率分布,从中选出得分最高的若干候选,进入下一轮。得分通常使用累积对数概率,这样可以把概率相乘转化为加法,避免浮点下溢,也便于比较不同长度序列。

束搜索的一个重要细节是长度惩罚。由于每个额外token都会让对数概率继续变小,纯累积分数天然偏向短序列,可能导致模型过早输出终止符。工程实现中一般会在最终排序时对分数除以序列长度,或者加入与长度相关的惩罚项。下面是一个简化版的束搜索实现,省略了具体的模型前向接口,只保留核心流程。

import math
def beam_search(model, input_ids, beam_width, max_steps):
    beams = [(input_ids, 0.0)]
    finished = []
    for _ in range(max_steps):
        candidates = []
        for seq, score in beams:
            if seq[-1] == model.eos_id:
                finished.append((seq, score))
                continue
            probs = model.step(seq)
            topk = probs.topk(beam_width)
            for token, p in zip(topk.indices, topk.values):
                new_seq = seq + [token]
                new_score = score + math.log(p + 1e-10)
                candidates.append((new_seq, new_score))
        candidates.sort(key=lambda x: x[1], reverse=True)
        beams = candidates[:beam_width]
        if len(beams) == 0:
            break
    finished.extend(beams)
    finished.sort(key=lambda item: item[1] / len(item[0]), reverse=True)
    return finished[0][0]

束搜索的主要优点是实现简单,计算过程高度并行。增加束宽可以显著降低早期选错分支的概率,但计算开销也会线性上升。它的局限也很明显:候选路径的评分完全依赖模型自身的概率输出,没有外部奖励信号。当模型在某个关键步骤给出了错误但置信度较高的预测时,束搜索仍然可能把所有候选都集中到错误方向上。此外,固定束宽意味着算力不会根据路径的潜在价值动态调整,这在需要长链条规划的推理任务中并不总是高效。

二、蒙特卡洛树搜索:把推理过程当作可回溯的决策树

蒙特卡洛树搜索最初在围棋等博弈问题中取得成功,它的核心思想是通过大量随机模拟来估计某个动作的长期期望收益。应用到推理模型时,状态可以定义为当前已经生成的部分推理步骤,动作则是下一步推理步骤或一小段token。算法通过选择、扩展、模拟、回传四个阶段反复迭代,逐步把算力集中到更有希望的分支上。

选择阶段从根节点出发,按照某种策略沿着树向下遍历。最常用的策略是UCT,它在利用已有平均收益和探索未知分支之间做平衡。节点访问次数越少,探索项越大,越容易被选中;而平均收益越高,利用项越大,也越容易被选中。下面的代码展示了UCT分数的计算方法。

import math
class MCTSNode:
    def __init__(self, state, parent=None):
        self.state = state
        self.parent = parent
        self.children = []
        self.visits = 0
        self.total_value = 0.0
    def uct_score(self, exploration=1.414):
        if self.visits == 0:
            return float('inf')
        exploit = self.total_value / self.visits
        explore = exploration * math.sqrt(
            math.log(self.parent.visits + 1) / self.visits
        )
        return exploit + explore

扩展阶段在到达叶子节点后生成新的子节点。与棋盘游戏不同,大语言模型的动作空间几乎是无限的,不可能枚举所有可能的下一个token。实际做法是从策略模型中采样若干个高概率候选,例如取top-k,或者使用温度采样生成多样化的下一步。模拟阶段从新节点出发快速生成一条完整路径,可以用小型模型、固定规则或直接调用critic网络来估计最终价值。回传阶段把模拟得到的奖励沿着路径更新到所有经过的节点上。

推理模型中使用蒙特卡洛树搜索的难点主要在于奖励设计。数学题可以判断最终答案是否正确,代码生成可以判断是否通过测试用例,但开放域问答和文本生成往往没有明确的自动验证手段。另一个难点是动作空间太大,即使只保留top-k候选,每次扩展仍然需要执行多次前向计算,总成本远高于束搜索。因此,工程实现中需要对节点数量、树深度和模拟次数设置严格预算。

三、混合策略与工程落地

Beam Search和蒙特卡洛树搜索并不是二选一的关系。工程上经常把两者组合起来:先用束搜索快速生成多条完整候选,再用测试用例或规则对这些候选打分,把得分作为蒙特卡洛树搜索的初始回传信号。这样既保留了束搜索的并行效率,又能利用外部奖励纠正模型概率估计的偏差。对于得分较高的候选路径,可以在其附近做小范围的树搜索,进一步探索变体。

维度Beam Search蒙特卡洛树搜索
搜索宽度固定动态分配
反馈信号模型概率外部奖励或价值网络
计算成本中等,易并行较高,需反复模拟
适用场景快速生成、低延迟长链条推理、有客观验证

在具体实现中,缓存键值对是降低混合搜索成本的关键。无论是束搜索的多个候选,还是蒙特卡洛树搜索的同层兄弟节点,它们的前缀往往相同。如果模型推理框架支持前缀缓存,可以避免重复计算共享token的注意力状态。另一个优化方向是批量扩展:把同一层所有待扩展节点的请求合并成一个batch,减少GPU调度开销。对于低概率分支,可以在进入搜索树之前直接剪枝,只保留概率高于阈值的候选。

奖励函数的设计也会显著影响搜索效果。以下是一个面向代码生成任务的简化奖励函数,它把测试通过率和长度惩罚结合起来,鼓励模型生成既正确又简洁的代码。

def reward_for_completion(completion, tests, length_penalty=0.001):
    passed = run_tests(completion, tests)
    if not passed:
        return 0.0
    return 1.0 / (1.0 + length_penalty * len(completion))

四、选择建议与调参实践

选择哪种搜索策略,首先要看任务有没有客观的验证信号。如果只是需要模型快速生成一段回答,没有测试用例或形式化验证,束搜索通常是更合理的选择,束宽设置从3到8一般能在质量和延迟之间取得平衡。如果任务允许事后验证,例如数学题可以检查答案、代码可以跑测试、规划可以用模拟器评估,那么蒙特卡洛树搜索或混合策略带来的准确率提升往往值得付出额外算力。

蒙特卡洛树搜索的调参比束搜索更复杂。模拟次数太少,价值估计方差很大;模拟次数太多,推理延迟会明显上升。探索系数在1到2之间比较常见,过大容易在低价值分支上浪费算力,过小则可能过早收敛到次优解。奖励值应当尽量归一化到固定区间,避免不同任务之间的数值差异影响搜索稳定性。温度采样参数也值得注意,较低的温度适合搜索阶段的确定性扩展,较高的温度适合前期多样性探索。

从实际效果看,Beam Search适合作为默认解码策略,蒙特卡洛树搜索适合处理那些失败成本较高、验证成本可控的推理任务。混合方案则提供了一个折中路径:用束搜索保证基础质量和速度,用树搜索对关键决策点做深度验证。随着推理模型越来越多地被用于数学证明、代码生成和智能体规划,这些搜索策略的落地方式还会继续细化,但核心原则不会变:根据反馈信号的强度和延迟预算,选择合适的搜索宽度与深度。

广告位 · 文章底部(后台-模块管理替换此区域)