Research map / Large language models / Agents, code and math
DS-Agent: Automated Data Science by Empowering Large Language Models with Case-Based Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DS-Agent, an LLM agent that applies case-based reasoning over Kaggle-derived human insight cases for an iterative development stage and a simplified retrieve-and-adapt deployment stage for automated machine learning pipelines.
Problem
LLM agents generate unreasonable experiment plans and fail to complete automated data science and ML model-building tasks reliably, and finetuning them is costly.
Approach
Integrates case-based reasoning into LLM agents: retrieves Kaggle-derived human insight cases, uses LLM-based ReviseRanker to reorder cases using execution feedback, and in deployment retrieves and adapts past successful solution code without iterative feedback, avoiding weight updates.
Key results
- DS-Agent with GPT-4 achieves 100% success rate across the 12 development tasks.
- In deployment, DS-Agent with GPT-4 reaches nearly 100% one pass rate and GPT-3.5 reaches 85%, versus best baselines at 60% and 56%.
- DS-Agent improves Mixtral-8x7B-Instruct one pass rate from 6% to 31%; deployment costs $0.0045 (GPT-3.5) and $0.135 (GPT-4) per run.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add CBR-based expert case retrieval and reuse (ablation) | LLM agent for data science · Development tasks (12) | average best rank: 3.41 -> 2.08 | Helped |
| add feedback-based case reranking (ablation) | DS-Agent · Development tasks (12) | average best rank: 2.58 -> 2.08 | Helped |
| adapt past successful solution cases instead of using textual human insights (ablation) | DS-Agent deployment stage · Deployment tasks | performance | Helped |
| add retrieved example cases to the prompt (ablation) | DS-Agent deployment stage · Deployment tasks (18) | performance | Helped |
| increase the number of example cases in context beyond one (ablation) | DS-Agent deployment stage · Deployment tasks | performance | Hurt |
| use retrieved similar cases rather than a random example case (ablation) | DS-Agent deployment stage · Deployment tasks (18) | average mean rank: 28.0 -> 12.7 | Helped |
| use retrieved similar cases rather than a random example case (ablation) | DS-Agent deployment stage · Deployment tasks (18) | average mean rank: 35.0 -> 17.1 | Helped |
| use case-based code adaptation instead of zero-shot code generation | Mixtral-8x7b-Instruct · Deployment tasks (18) | one-pass rate: 6% -> 31% | Helped |
| use DS-Agent rather than AutoGluon | GPT-4 data science model · wild-blueberry-yield | MAE: 305 -> 182 | Helped |
| use DS-Agent rather than AutoGluon | GPT-4 data science model · spaceship-titanic | accuracy: 0.8044 -> 0.8062 | Helped |
| use DS-Agent rather than AutoGluon | GPT-4 data science model · enzyme-substrate | AUROC: 0.8683 -> 0.9761 | Helped |
| use DS-Agent rather than AutoGluon | GPT-4 data science model · media-campaign-cost | RMLSE: 0.2707 -> 0.2858 | Hurt |
Surprising findings
- Adding more retrieved example cases in the deployment context degrades performance, contrary to typical few-shot learning expectations.
- Learning from past successful solution code outperforms learning from textual human insights in deployment.
Limitations
- Development stage is costly ($1.60 per run with GPT-4) and relies on closed-source API models.
- Deployment performance of the open-source Mixtral model remains weaker than GPT-3.5 and GPT-4.
- Evaluated on 30 tasks across three modalities; comparison with AutoML is limited to four tabular tasks.
Experimental setup
Datasets: Kaggle competition technical reports and notebooks (12 competitions) 30 data science tasks (12 development, 18 deployment)
Benchmarks: DS-Agent 30-task suite (Kaggle and research datasets)
Metrics: success rate one pass rate mean rank best rank monetary cost
Baselines: ResearchAgent Zero-shot One-shot AutoGluon w/o ReviseRank w/o CBR
Builds on: Reflexion ResearchAgent retrieval-augmented generation case-based reasoning
Closest papers by meaning
- DSBench: How Far Are Data Science Agents from Becoming Data Science Experts? (ICLR 2025 · similarity 0.90)
- Distilling LLM Agent into Small Models with Retrieval and Code Tools (NeurIPS 2025 · similarity 0.87)
- AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML (ICML 2025 · similarity 0.87)
- AgentBench: Evaluating LLMs as Agents (ICLR 2024 · similarity 0.86)
- MAS-GPT: Training LLMs to Build LLM-based Multi-Agent Systems (ICML 2025 · similarity 0.86)
- Star-Agents: Automatic Data Optimization with LLM Agents for Instruction Tuning (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.