更宽还是更深？通过自适应分支树搜索扩展 LLM 推理时计算

ArXiv ID: 2503.04412
作者: Yuichi Inoue, Kou Misaki, Yuki Imajuku, So Kuroki, Taishi Nakamura, Takuya Akiba
机构: Preferred Networks
发布日期: 2025-03-06
接收: ICLR 2025 Workshop, NeurIPS 2025 Spotlight

---

摘要

在 LLM 推理时扩展（test-time scaling）中，一个核心问题是：应该探索更多不同的回答（更宽），还是深入改进已有回答（更深）？本文提出的 AB-MCTS（Adaptive Branching MCTS） 框架通过自适应地平衡这两种策略，在编码任务上显著优于重复采样和标准 MCTS 方法。

---

核心问题

推理时扩展的两难选择

更宽策略（Width）：
- 生成 N 个独立回答
- 优点：多样性高，覆盖更多解空间
- 缺点：每个回答质量可能不高
- 代表方法：重复采样（Repeated Sampling）

更深策略（Depth）：
- 对单个回答进行迭代改进
- 优点：可以得到高质量解
- 缺点：可能陷入局部最优
- 代表方法：标准 MCTS、Self-Refine

关键洞察

问题类型决定策略：

• 需要探索的问题（如开放性问题）→ 更宽更好
• 需要精确的问题（如数学证明）→ 更深更好
• 大多数实际问题 → 需要动态平衡

---

AB-MCTS 方法

整体架构

┌─────────────────────────────────────────────────────────┐
│                   AB-MCTS Controller                     │
│  ┌──────────────┐        ┌──────────────┐              │
│  │  宽度决策器   │        │  深度决策器   │              │
│  │  (Expand)    │        │  (Improve)   │              │
│  └──────────────┘        └──────────────┘              │
│           │                      │                       │
│           └──────────┬───────────┘                       │
│                      ▼                                   │
│           ┌─────────────────┐                           │
│           │  自适应选择器    │                           │
│           │  (基于反馈信号)  │                           │
│           └─────────────────┘                           │
└─────────────────────────────────────────────────────────┘
                          │
                          ▼
┌─────────────────────────────────────────────────────────┐
│                    树搜索过程                            │
│                                                         │
│     Root                                                │
│    /  |  \                                               │
│   ○   ○   ○  ← 宽度扩展 (生成新候选)                      │
│   |       |                                              │
│   ○       ○  ← 深度改进 (迭代优化)                        │
│                                                         │
└─────────────────────────────────────────────────────────┘

算法流程

def ab_mcts(problem, llm, feedback_fn, budget):
    """
    AB-MCTS 主算法

    Args:
        problem: 待解决问题
        llm: 语言模型
        feedback_fn: 外部反馈函数（如测试用例、验证器）
        budget: 计算预算（最大迭代次数）
    """
    # 初始化树
    root = Node(problem)
    tree = Tree(root)

    for iteration in range(budget):
        # 步骤 1: 选择节点
        node = tree.select()

        # 步骤 2: 自适应决策 - 宽度还是深度？
        if should_expand(node, tree, feedback_fn):
            # 宽度扩展：生成新的候选回答
            new_responses = llm.generate_multiple(
                prompt=node.problem,
                num_samples=branch_factor
            )
            for resp in new_responses:
                tree.add_child(node, resp)
        else:
            # 深度改进：改进现有回答
            improved = llm.improve(
                current_answer=node.answer,
                feedback=feedback_fn(node.answer)
            )
            node.update(improved)

        # 步骤 3: 获取反馈并更新
        feedback = feedback_fn(node.answer)
        node.update_score(feedback)

        # 步骤 4: 反向传播
        tree.backpropagate(node, feedback)

    # 返回最佳回答
    return tree.get_best_answer()


def should_expand(node, tree, feedback_fn):
    """
    自适应决策：是否应该扩展新分支

    决策依据：
    1. 当前节点的反馈分数
    2. 兄弟节点的数量
    3. 树的探索程度
    """
    # 如果当前节点分数低，说明需要探索新方向
    if node.score < threshold:
        return True

    # 如果探索不足，优先扩展
    if tree.exploitation_ratio < target_ratio:
        return True

    # 否则进行深度改进
    return False

自适应机制

决策因子：

因子	描述	影响
当前分数	节点的回答质量	低分→扩展
分支数	已有候选数量	少→扩展
改进幅度	上次改进的收益	小→扩展
预算剩余	剩余计算资源	多→探索

决策公式：

P(expand) = σ(w1·(1-score) + w2·(1-branches/max) + w3·(1-improvement) + w4·budget_ratio)

---

实验设置

基准任务

编码任务：

• LeetCode 中等/困难题目
• Codeforces 竞赛题
• AtCoder 题目

反馈机制：

• 测试用例通过率
• 代码执行结果
• 编译成功/失败

对比方法

方法	类型	描述
Repeated Sampling	纯宽度	生成 N 个独立回答，选最佳
Standard MCTS	平衡	标准蒙特卡洛树搜索
Self-Refine	纯深度	迭代改进单个回答
Chain of Thought	单次	思维链推理

评估指标

• 通过率：正确解决题目的比例
• 改进率：相比单次生成的提升
• 效率：单位计算量的性能
• 多样性：生成回答的差异化程度

---

实验结果

主要结果

方法	LeetCode Medium	LeetCode Hard	Codeforces
Single Pass	45.2%	23.1%	15.3%
Repeated Sampling (×8)	58.7%	31.4%	22.1%
Standard MCTS	55.3%	29.8%	20.5%
Self-Refine (×8)	52.1%	28.6%	18.9%
AB-MCTS (×8)	64.3%	38.2%	27.8%

关键发现：

1. AB-MCTS 在所有基准上均优于基线方法
2. 相比纯宽度策略，提升 5-7 个百分点
3. 相比标准 MCTS，提升 8-9 个百分点

消融实验

自适应机制的有效性：

变体	LeetCode Medium	说明
AB-MCTS (完整)	64.3%	完整模型
- 固定宽度	58.1%	固定分支因子
- 固定深度	55.9%	固定改进次数
- 随机决策	52.4%	随机宽度/深度选择

计算效率分析

性能 vs 计算量：

通过  │
率    │           ● AB-MCTS
      │        ●
      │     ●
      │  ●           ▲ Repeated Sampling
      │            ▲
      │         ▲
      │      ▲
      └─────────────────────
         计算预算 (迭代次数)

AB-MCTS 在相同计算预算下始终取得更高性能，或在达到相同性能时所需计算更少。

案例研究

成功案例：

问题：实现一个支持 O(1) 插入、删除、随机返回的数据结构

AB-MCTS 的搜索过程：
1. [宽度] 生成初始方案：使用列表 → 测试失败（删除不是 O(1)）
2. [宽度] 生成新方案：使用字典 → 测试失败（随机返回不是 O(1)）
3. [深度] 改进方案 2：字典 + 列表混合 → 测试通过 ✓

关键：在第 2 步后，算法判断继续改进方案 2 比探索新方向更有价值

失败分析：

• 需要创造性洞察的问题（如数学证明）表现较差
• 反馈信号稀疏时决策质量下降
• 极小概率事件可能导致错误决策

---

讨论

优势

1. 自适应性：根据问题特征动态调整策略
2. 原则性：基于 UCT 公式的理论保证
3. 通用性：适用于各种有反馈的任务
4. 高效性：在相同预算下取得更好结果

局限性

1. 依赖反馈：需要可靠的外部反馈信号
2. 计算开销：决策机制本身有额外成本
3. 参数敏感：某些超参数需要调优

适用场景

适合 AB-MCTS 的任务：

• 有明确验证标准的任务（编码、数学计算）
• 解空间结构化、可探索的任务
• 迭代改进可行的任务

不太适合的任务：

• 主观性强的创意写作
• 反馈延迟或稀疏的任务
• 单次通过即可的任务

---

与相关工作的关系

推理时扩展谱系

推理时扩展方法
│
├── 采样基方法
│   ├── Repeated Sampling（纯宽度）
│   ├── Best-of-N
│   └── Consensus-based
│
├── 树搜索方法
│   ├── Standard MCTS
│   ├── AlphaCode-style
│   └── AB-MCTS (本文)
│
└── 迭代改进方法
    ├── Self-Refine
    ├── Iterative Correction
    └── Process Supervision

与 O1 等系统的比较

特性	OpenAI o1	AB-MCTS
推理方式	内部思考链	外部树搜索
计算分配	隐式学习	显式自适应
透明度	低	高
可解释性	思考过程不可见	树结构可视化

---

实践建议

实现要点

1. 反馈设计：

   • 确保反馈信号可靠、快速
   • 考虑部分正确的细粒度评分

2. 参数调优：

   • 初始探索率：0.3-0.5
   • 深度/宽度权衡因子：根据任务调整
   • 预算分配：预留 20% 用于最终验证

3. 工程优化：

   • 缓存已评估的回答
   • 并行评估多个叶节点
   • 早停机制：达到满意解即可停止

代码示例

from ab_mcts import ABMCTS, LLMWrapper, FeedbackProvider

# 配置
llm = LLMWrapper(model="claude-sonnet-4")
feedback = FeedbackProvider(type="code_execution")

solver = ABMCTS(
    llm=llm,
    feedback=feedback,
    max_iterations=50,
    exploration_weight=0.4,
    branch_factor=3
)

# 解决问题
result = solver.solve("LeetCode 146: LRU Cache 实现")
print(f"解决成功：{result.success}")
print(f"迭代次数：{result.iterations}")
print(f"最佳方案：{result.code}")

---

总结

AB-MCTS 通过自适应地平衡”更宽”探索和”更深”改进，为 LLM 推理时计算提供了一个原则性的解决方案。实验证明，这种方法在编码任务上显著优于固定策略，为实现更高效的推理时扩展开辟了新方向。

核心贡献：

1. 提出了 AB-MCTS 框架，自适应平衡宽度与深度
2. 在编码基准上取得 SOTA 结果
3. 提供了关于推理时计算分配的新洞察

未来方向：

• 扩展到多模态任务
• 学习更复杂的决策策略
• 与模型内部推理机制结合

---

资源

• arXiv 论文
• GitHub 代码
• ICLR Workshop 演讲视频

---

评分: 4.2/5.0

推荐度: ⭐⭐⭐⭐ 适合需要推理时扩展的开发者参考