跳转至

Jupiter: Enhancing LLM Data Analysis Capabilities via Notebook and Inference-Time Value-Guided Search

会议: AAAI 2026
arXiv: 2509.09245
代码: https://github.com/microsoft/Jupiter
领域: LLM推理
关键词: 数据分析, Jupyter Notebook, 值引导搜索, MCTS, 测试时计算

一句话总结

构建NbQA数据集（从真实Jupyter Notebook提取3.8万task-solution对）+ 提出Jupiter框架（将数据分析建模为状态级搜索问题，用值模型引导PUCT搜索），使Qwen2.5-14B在InfiAgent-DABench上达86.38%超越GPT-4o(85.99%)，Qwen2.5-7B在DSBench上从63.51%提升至89.19%。

研究背景与动机

  1. 领域现状：LLM在数据分析任务上的应用日益广泛，但现有方法要么关注孤立的分析阶段，要么依赖专有模型（如GPT-4o）和复杂agent架构（AutoGen、TaskWeaver），开源模型在多步数据分析上表现不佳。

  2. 现有痛点：

  3. 数据不足：缺乏大规模、高质量的数据分析训练数据。现有数据集多由LLM合成，质量和多样性受限
  4. 推理能力弱：数据分析需要多步推理和工具使用，LLM在复杂分析链中容易出错且无法自我纠正

  5. Agent系统依赖闭源模型：AutoGen等框架效果好但依赖GPT-4级别的模型

  6. 核心矛盾：开源小模型（7B-14B）在数据分析上远逊于闭源大模型，缺乏高质量训练数据和有效的推理时搜索策略。

  7. 本文要解决什么？ (1)从真实Notebook构建高质量训练数据 (2)设计推理时值引导搜索提升多步数据分析准确率。

  8. 切入角度：观察到Jupyter Notebook天然包含问题-代码-执行结果的多步结构，是理想的数据分析训练来源；同时MCTS+值模型在推理任务上的成功可以迁移到数据分析场景。

  9. 核心idea一句话：真实Notebook数据 + 值模型引导的推理时搜索 = 开源小模型匹敌GPT-4o的数据分析能力。

方法详解

整体框架

两大组件：(1) NbQA数据集构建——从1.6M真实Jupyter Notebook中提取高质量task-solution对；(2) Jupiter推理框架——将数据分析建模为搜索树，用MCTS收集轨迹训练值模型，推理时用值模型引导搜索。

关键设计

  1. NbQA数据集构建Pipeline:
  2. 做什么：从真实Jupyter Notebook中提取标准化的数据分析task-solution对
  3. 核心思路：粗过滤（有效结构、成功执行、足够复杂度）→ GPT-4o mini质量评分(1-5) → GPT-4o精细提取（标准化格式）→ GPT-4o mini自动审核。最终得到38,635个任务对，6,845个含完整数据依赖
  4. 覆盖8类任务：统计摘要、分布分析、相关分析、异常检测、数据预处理、特征工程、机器学习、可视化

  5. 设计动机：真实Notebook解决方案比LLM合成方案更多样、更实际，直接反映数据科学家的实际工作模式

  6. Jupiter搜索框架:

  7. 做什么：将数据分析建模为状态级搜索问题，构建搜索树进行探索
  8. 核心思路：根节点=原始问题，节点=Notebook状态（累积的thought-action+执行结果），叶节点=错误/候选答案。使用PUCT公式选择节点：\(\text{PUCT}(s') = Q(s') + c_{\text{puct}} \cdot P(s') \cdot \frac{\sqrt{N_{\text{parent}}}}{1+N(s')}\)
  9. 训练阶段用MCTS（\(c_{\text{puct}}>0\)）收集多样轨迹；推理阶段关闭探索项（\(c_{\text{puct}}=0\)），仅用值模型引导

  10. 设计动机：数据分析是顺序决策问题，搜索树能系统地探索不同分析路径

  11. 值模型训练:

  12. 做什么：训练回归头预测每个Notebook状态的Q值
  13. 核心思路：在fine-tuned LLM上附加回归头，输出[-1,1]标量。用MCTS收集的轨迹中的归一化Q值作为监督信号。3个epoch，batch size 4，学习率1e-4
  14. 设计动机：值模型将MCTS中学到的"哪条路径更可能成功"的知识压缩为一个评估函数，推理时无需完整MCTS

损失函数 / 训练策略

  • SFT阶段：用NbQA数据对基础模型做supervised fine-tuning
  • 值模型训练：回归损失，输入为完整多轮对话，输出为MCTS轨迹的归一化Q值
  • 尝试了GRPO-based RL但效果不如SFT（稀疏奖励+规则正确性检查的问题）

实验关键数据

主实验

InfiAgent-DABench（40迭代）：

模型 框架 准确率
GPT-4o TaskWeaver 85.99%
GPT-4o ReAct 81.32%
Qwen2.5-72B ReAct 75.88%
Qwen2.5-14B (SFT) Jupiter 86.38%
Qwen2.5-7B (SFT) Jupiter 77.82%

DSBench数据建模：

模型 成功率
GPT-4 (AutoGen) 87.84%
Qwen2.5-7B (Jupiter) 89.19%
Qwen2.5-14B (Jupiter) 98.65%

消融实验

配置 Qwen2.5-7B Acc Qwen2.5-14B Acc
w/o VM, w/o ExpTerm 70.04% 79.38%
w/o VM, with ExpTerm 68.87% 75.88%
with VM, with ExpTerm 68.87% 74.71%
with VM, w/o ExpTerm 77.82% 86.38%

关键发现

  • 值模型至关重要：带vs不带值模型差7-8个点
  • 探索项反而有害：关闭探索项后准确率提升8-11个百分点，因为数据分析的搜索空间稀疏，盲目探索不如精准引导
  • NbQA数据质量好：仅SFT就使Mistral-7B提升56.81%（2.33%→59.14%），说明真实Notebook数据极有价值
  • 泛化到数学推理：在AIME数学竞赛上，Qwen2.5-7B从0%提升到33.3%（OR）/20%（voting）

亮点与洞察

  • 真实Notebook > 合成数据：从1.6M真实Jupyter Notebook中精炼出的训练数据远优于LLM合成数据，因为它反映了数据科学家的真实工作流程
  • 推理时关闭探索：与AlphaGo不同数据分析任务搜索空间稀疏，有好的值模型后应做贪心搜索而非探索
  • 开源小模型超越闭源大模型：14B模型超越GPT-4o，证明了"好数据+好搜索"可以弥补模型规模差距

局限性 / 可改进方向

  • RL训练失败：GRPO-based RL不如SFT，可能因为正确性检查规则太粗（部分正确不给分）
  • 值模型泛化有限：值模型在训练域(数据分析)上有效，但跨域（如数学推理）效果存疑
  • 搜索开销大：40次迭代的搜索在推理时成本高，需要更高效的搜索策略
  • 数据分析任务定义窄：仅覆盖结构化数据分析，未包含非结构化数据处理、深度学习等

相关工作与启发

  • vs Data Interpreter: Data Interpreter是复杂agent系统，Jupiter是简洁的搜索+值模型框架，在更小模型上达到更好效果
  • vs AlphaMath/MCTS-RL: 数学推理中MCTS用探索发现新路径，数据分析中关闭探索反而更好——任务性质决定搜索策略
  • vs ReAct: ReAct是无搜索的单轨推理，Jupiter通过搜索树系统地尝试多条分析路径

评分

  • 新颖性: ⭐⭐⭐⭐ Notebook数据提取+值引导搜索的组合是实用创新，但各组件并非全新
  • 实验充分度: ⭐⭐⭐⭐⭐ 两个基准+多个模型大小+详细消融+超参数分析+跨域泛化测试
  • 写作质量: ⭐⭐⭐⭐ 方法描述清晰，数据构建pipeline详细，但论文较长
  • 价值: ⭐⭐⭐⭐⭐ NbQA数据集和Jupiter框架均开源，对数据分析自动化社区有重要贡献