Skip to content

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

🕒 Published (v1): 2025-09-11 08:27 UTC · Source: Arxiv · Venue: AAAI 2026 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Jupiter is a framework that casts multi-step data analysis in Jupyter notebooks as a tree-search problem, using Monte Carlo Tree Search (MCTS) during training to collect trajectories and a learned value model to guide inference-time search. Paired with NbQA—a 38,635-pair dataset extracted from 1.6M real-world Jupyter notebooks—fine-tuned 7B and 14B models reach 77.82% and 86.38% on InfiAgent-DABench, matching or exceeding GPT-4o-backed agent frameworks.

Problem

LLMs struggle with multi-step data analysis: existing agent frameworks either focus on single pipeline stages, rely on proprietary models, or use ad-hoc decomposition without principled search. There is also no large-scale, tool-grounded dataset of authentic multi-step data analysis trajectories drawn from real notebooks with executable dependencies.

Method

NbQA construction: GitHub API crawls 1.6M .ipynb notebooks and 3.2M data files. Coarse filters remove malformed, unexecuted, or contaminated notebooks; GPT-4o mini scores quality; GPT-4o extracts 1–3 task–solution pairs per notebook (thought–code–output blocks) across eight task categories, annotated with explicit constraints and @answer labels. Final dataset: 38,635 pairs, 6,845 with fully interactive data dependencies.

Jupiter framework: Each notebook state \(s\) is a node in a search tree; edges correspond to generating and executing one code cell. Node selection uses the PUCT score: $\(\text{PUCT}(s') = Q(s') + c_{\text{puct}} \cdot P(s') \cdot \frac{\sqrt{N_{\text{parent}}(s')}}{1 + N(s')}\)$ During trajectory collection (\(c_{\text{puct}} > 0\)), MCTS explores the tree and backpropagates \(\pm 1\) rewards. Q-values from these simulations supervise a value model—an LLM with a regression head trained via MSE to output a scalar in \([-1, 1]\). At inference, \(c_{\text{puct}} = 0\) removes exploration, relying entirely on value estimates and visit counts to focus search on high-quality branches. Final answer is chosen by majority vote over all candidate answer nodes.

Key Contributions

  • Automated pipeline for mining high-quality, executable multi-step data analysis task–solution pairs from real Jupyter notebooks at scale.
  • NbQA: 38,635 standardized task–solution pairs across 8 analysis types, supporting both SFT and value model training.
  • Jupiter framework: MCTS-based trajectory collection + value model training + inference-time value-guided search with exploration term disabled.
  • Demonstration that a value model trained on NbQA transfers to out-of-distribution tasks (DSBench Kaggle modeling, AIME math competition) without task-specific adaptation.

Results

  • InfiAgent-DABench (SFT only): Qwen2.5-7B-Instruct: +24.12 pp (43.97→68.09%); Qwen2.5-14B-Instruct: +7.39 pp (69.65→77.04%); Mistral-7B: +56.81 pp; Llama-3.1-8B: +21.40 pp.
  • InfiAgent-DABench (Jupiter, 40 iters, with VM, w/o ExpTerm): Qwen2.5-7B: 77.82%, Qwen2.5-14B: 86.38% — surpasses Taskweaver+GPT-4o (85.99%) and all other baselines.
  • DSBench (50 iters, vanilla model + VM): Qwen2.5-14B: 98.65% task completion; Qwen2.5-7B: 89.19% — both exceed GPT-4-based frameworks (Code Interpreter, AutoGen).
  • AIME 2025 (out-of-domain math): After SFT, Jupiter+VM raises voting accuracy to 13.3% (vs. 0% tool-use baseline); pass@k rises to 26.7%.
  • Removing the exploration term consistently outperforms keeping it, confirming the value model alone is sufficient for efficient search.

Limitations

  • Value model training requires 6,845 fully interactive (data-file-complete) notebook samples; broader NbQA pairs without interactive dependencies cannot be used for MCTS trajectory collection.
  • NbQA excludes neural-network and GPU-based notebooks, limiting coverage of deep-learning-centric data science workflows.
  • AIME results remain modest (13.3% voting accuracy), suggesting the math tool-use gains are real but not yet competitive with math-specialized models.
  • Inference-time search with 40 iterations at \(K=3\) expansions per step is computationally expensive relative to single-pass inference; no latency/cost analysis is reported.

Relevance to Harnesses / Meta-Harnesses

Jupiter is a concrete example of an inference-time meta-harness: it wraps a base LLM in a structured search loop (MCTS), executes code cells in a live Jupyter kernel, collects and routes outputs back into the tree state, and uses a learned value model to decide which branch to expand next—exactly the kind of orchestration layer that harness/meta-harness research addresses. The NbQA pipeline itself is a data-generation harness: an automated chain of crawling, filtering, LLM-based scoring, and extraction that produces training signal without human annotation. The finding that disabling the exploration term (pure value-guided exploitation) outperforms UCB-style exploration is directly relevant to harness design choices around when to explore vs. commit—analogous to routing and retry policies in orchestration frameworks.