PaperlayerView exact revision on arXiv

Source-linked AI summary

DS-Agent: Automated Data Science by Empowering Large Language Models with Case-Based Reasoning

Siyuan Guo, Cheng Deng, Ying Wen, Hechang Chen, Yi Chang, Jun Wang

arXiv:2402.17453v5cs.LG

TL;DR

LLM agents struggle to complete data-science tasks, with unreasonable experiment planning identified as a weakness. DS-Agent combines case-based reasoning with iterative development and low-resource reuse of successful solutions; across 30 tasks, it shows strongest performance in both stages, including 100% GPT-4 success on 12 development tasks.

  • Problem

    Existing LLM agents struggle to achieve high task-completion rates in data science, with unreasonable planning and hallucinations identified as contributing weaknesses.

    Source: §1 · p. 1

  • Method

    DS-Agent uses case-based reasoning to retrieve Kaggle insights and iteratively refine plans from execution feedback, then reuses successful development-stage solutions for low-resource deployment code generation.

    Source: §1 · p. 2

  • Results

    Across 30 data-science tasks, DS-Agent showed strongest performance in both stages, with GPT-4 achieving 100% success across 12 development tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Reusing successful cases supports low-resource deployment on unseen tasks in the same distribution while reducing demands on foundational LLM capabilities.

    Source: §1 · p. 2

    Generated code may be inefficient, exploitable, or harmful, so users should review it before execution.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

In this work, we investigate the potential of large language models (LLMs) based agents to automate data science tasks, with the goal of comprehending task requirements, then building and training the best-fit machine learning models. Despite their widespread success, existing LLM agents are hindered by generating unreasonable experiment plans within this scenario. To this end, we present DS-Agent, a novel automatic framework that harnesses LLM agent and case-based reasoning (CBR). In the development stage, DS-Agent follows the CBR framework to structure an automatic iteration pipeline, which can flexibly capitalize on the expert knowledge from Kaggle, and facilitate consistent performance improvement through the feedback mechanism. Moreover, DS-Agent implements a low-resource deployment stage with a simplified CBR paradigm to adapt past successful solutions from the development stage for direct code generation, significantly reducing the demand on foundational capabilities of LLMs. Empirically, DS-Agent with GPT-4 achieves 100\% success rate in the development stage, while attaining 36\% improvement on average one pass rate across alternative LLMs in the deployment stage. In both stages, DS-Agent achieves the best rank in performance, costing \$1.60 and \$0.13 per run with GPT-4, respectively. Our data and code are open-sourced at https://github.com/guosyjlu/DS-Agent.

1. Introduction

DS-Agent combines LLMs with case-based reasoning (CBR) to automate model-centric data science, using Kaggle expertise and execution feedback during development and past successful solutions for low-resource deployment.

1. Introduction

  • Existing LLM agents struggle to complete data-science tasks because they produce unreasonable plans and hallucinate, while finetuning requires costly task-execution feedback and substantial computation.The cited work reports these limitations even for agents using GPT-4.Source: §1 · p. 1
  • Kaggle’s repository of data-science reports and code motivates CBR, which retrieves similar past problems, reuses and revises their solutions, evaluates outcomes, and retains successful cases.This lets agents draw on expert insights and execution feedback while conserving sample and computation resources.Source: §1 · pp. 1–2
  • In development, DS-Agent retrieves Kaggle insights to ground experiment plans, revises cases and plans using execution feedback, and retains successful solutions instead of updating model parameters.The feedback-driven iterations are described as improving performance consistently.Source: §1 · p. 2
  • For low-resource deployment, DS-Agent retrieves successful development-stage solutions and adapts them with minor modifications to generate code for unseen tasks in the same distribution, without execution-feedback revision.This simplified CBR approach reduces the demands on LLM foundational capabilities.Source: §1 · p. 2
  • Across 12 development tasks, GPT-4 DS-Agent achieved 100% success; on 18 deployment tasks, GPT-3.5 and GPT-4 achieved 85% and 99% one pass rates versus 56% and 60% for the best baseline.For Mixtral-8x7b-Instruct, the one pass rate rose from 6% to 31%.Source: §1 · p. 2
  • DS-Agent ranked first with GPT-4 and second with GPT-3.5 in both stages, costing $1.60 and $0.06 per run in the standard scenario and $0.135 and $0.0045 in low-resource scenarios.The per-run costs are reported in GPT-4 then GPT-3.5 order for each scenario.Source: §1 · p. 2

2. Preliminary

CBR-based LLMs iteratively retrieve similar cases, generate solutions, evaluate them, and revise them using feedback. Unlike RAG, the framework can adapt retrieved cases and retain effective solutions for future reuse.

CBR based LLMs

  • CBR-based LLMs retrieve similar problems, generate and evaluate solutions, revise them as needed, and retain the best-performing solution for future reuse.The framework uses a retriever, an LLM, and an evaluator that produces feedback.Source: §2 · p. 2
  • Evaluating the current solution distribution produces feedback that revises the subsequent distribution, forming an iteration loop intended to improve performance consistently.The solution distribution is conditioned on the previous step’s execution feedback and retrieved case.Source: §2 · p. 2

Comparison with Retrieval-Augmented Generation

  • Unlike RAG, which uses a retriever and an LLM, CBR-based LLMs adjust retrieved cases and revise solutions in response to evaluation feedback.Retaining good solutions in the database provides a flexible learning mechanism for consistent performance improvement.Source: §2 · p. 3

3. The DS-Agent

DS-Agent uses case-based reasoning (CBR) in two stages: an iterative development pipeline that improves experiment plans using expert cases and execution feedback, and a lower-resource deployment pipeline that adapts successful past solutions into code.

3. The DS-Agent

  • Across both stages, DS-Agent interprets a task description, trains an ML model on training data, evaluates it on validation data, and reports test-set performance.The framework separates tasks into development and deployment sets.Source: §3 · p. 3

3.1. Development Stage: Automatic Iteration Pipeline

  • In development, CBR uses expert cases to guide experiment plans, revises case rankings and plans from execution feedback, and retains successful solutions without resource-intensive LLM fine-tuning.The paper reports a consistent trend of performance improvement across iteration steps.Source: §3.1 · pp. 3–5

3.2. Deployment Stage: Learning from Past Cases

  • In deployment, DS-Agent retrieves a similar successful case from its agent case bank and adapts its code, reducing LLM reasoning and coding demands for unseen tasks in the same task distribution.This simplified CBR approach supports implementation with open-source LLMs.Source: §3.2 · p. 5

4. Experiments

Experiments evaluate DS-Agent across development and deployment tasks using completion, model-performance, and monetary-cost metrics. Results support the CBR approach in both stages and examine how feedback, prior cases, and the number of contextual examples affect performance.

4.1. Experiment Setting

  • The evaluation covered 30 tasks spanning text, time-series, and tabular data and regression and classification, split into 12 development and 18 deployment tasks.Completion was measured by success rate in development and one-pass rate in deployment; model performance by mean and best rank; and resource cost by money consumed.Source: §4.1 · p. 5

4.2.1. MAIN RESULTS

  • In development, DS-Agent with GPT-4 achieved the highest success rate (100% over all tasks) and best performance on 9 of 12 tasks, outperforming other agents in mean and best rank; GPT-3.5 had the second-best average result in both.For both GPT-3.5 and GPT-4, DS-Agent’s average best mean rank improved over increasing iteration steps.Source: §4.2.1 · p. 6

4.2.2. ABLATION STUDY

  • Without feedback-based case revision, performance deteriorated; without CBR, performance was worst among the three agents.The no-CBR variant generated plans without human insights, while CBR enabled use of Kaggle expert knowledge.Source: §4.2.2 · pp. 6–7, Table 2 · p. 6

4.3.1. MAIN RESULTS

  • On 18 deployment tasks, DS-Agent with GPT-4 attained the highest mean rank among nine agents, while its open-source version outperformed or competed with open-source baselines on 13 tasks.DS-Agent with GPT-3.5 ranked second in mean rank, surpassing GPT-4 baselines; the open-source model remained weaker than GPT-3.5 and GPT-4 agents overall.Source: §4.3.1 · p. 7, Table 3 · p. 7
  • Deployment cost per run was $0.0045 with GPT-3.5 and $0.1350 with GPT-4, over 90% below development-stage costs.The deployment stage is designed to solve tasks with minimal resources after development explores effective model designs.Source: §4.3.1 · pp. 7–8, Table 4 · p. 8

4.3.2. FURTHER ANALYSES

  • Past successful task-solution cases outperformed textual human insights, while adding more than one contextual case degraded performance through interference.With no example cases, DS-Agent reverted to zero-shot and performed worst among the tested settings.Source: §4.3.2 · p. 8, Figure 6 · p. 8

5. Related Work

The related work situates DS-Agent among LLM agents for data science and presents case-based reasoning (CBR) as a way to reuse and iteratively improve solutions. It distinguishes DS-Agent’s focus on building and training machine-learning models from other agents’ tasks.

LLM Agent

  • DS-Agent focuses on automatically building and training ML models, unlike prior LLM agents for feature engineering, hyperparameter tuning, ML-library use, AI research, or data operations.Earlier work also discussed LLMs as conversational agents in data science workflows (Mahadi Hassan et al., 2023).Source: §5 · p. 8
  • Data Interpreter optimizes LLM-agent workflows for data science, while its techniques and DS-Agent’s are complementary (Hong et al., 2024a).The authors suggest future work could enhance DS-Agent with Data Interpreter or integrate CBR into Data Interpreter.Source: §5 · p. 8

Case-Based Reasoning

6. Conclusion

DS-Agent applies case-based reasoning (CBR) to data science tasks through an iterative development stage and a low-resource deployment stage that reuses successful solutions. The paper reports that extensive experiments demonstrate its effectiveness.

6. Conclusion

  • DS-Agent combines an LLM agent with case-based reasoning to solve data science tasks.The framework has distinct development and deployment stages.Source: §6 · p. 9
  • In development, DS-Agent retrieves Kaggle insights to plan experiments, then revises retrieved cases and plans using execution feedback.Source: §6 · p. 9
  • In deployment, DS-Agent uses simplified CBR to retrieve and reuse successful solution cases from development in a low-resource setting.Source: §6 · p. 9
  • Extensive experiments demonstrate DS-Agent’s effectiveness for data science tasks.Source: §6 · p. 9

Impact Statement · Appendix · A. Pseudo-code of DS-Agent

The impact statement describes DS-Agent as assistive while flagging code and privacy risks; the appendix pseudocode specifies iterative development and case-based deployment.

Impact Statement

  • DS-Agent is intended to assist data scientists with complex work, not replace them, and natural-language interaction may broaden access to data insights.The described higher-level work includes task formulation, visualization, data cleaning and curation, prediction engineering, and summarizing and recommending results.Source: Impact Statement · p. 9
  • DS-Agent could generate inefficient, exploitable, or harmful code; users should review generated code and the paper recommends running it in Docker.Such issues were not observed in the experiments; Docker is recommended to isolate the host file system.Source: Impact Statement · p. 9
  • Local operation avoids uploading sensitive data, but API-based LLMs such as GPT-3.5 or GPT-4 may transmit data externally, so prompts should be checked for disclosures.Source: Impact Statement · p. 9

A. Pseudo-code of DS-Agent

  • In development, DS-Agent retrieves similar human-insight cases, reranks them for the task and current log, selects one, then plans, generates, and executes code.The pseudocode initializes an empty agent experience case bank and uses separate ranking, planning, and programming agents.Source: A. Pseudo-code of DS-Agent · p. 13
  • Development continues through bounded debugging, logs each experiment, and stores task descriptions and code when test-set performance improves.The debugger corrects code and reruns it while errors remain and the debugging count is below N.Source: A. Pseudo-code of DS-Agent · p. 13
  • In deployment, DS-Agent retrieves the most similar past case from its agent case bank and uses an adaptation agent to generate code for the new task.Source: A. Pseudo-code of DS-Agent · p. 13

B. Experimental Details · B.1. Task Selection

DS-Agent’s evaluation uses 30 representative data-science tasks spanning three modalities and two fundamental ML task types, with task-specific metrics and natural-language prompts. Most datasets postdate September 2021, and each task includes a random-guess baseline as an initial reference.

B.1. Task Selection

  • The evaluation covers 30 representative tasks across three data modalities and two fundamental ML task types.Detailed task descriptions are presented in Table 5.Source: §B.1 · p. 13, Table 5 · p. 15
  • Most selected datasets were released after September 2021, so they were not part of LLM pretraining corpora.Source: §B.1 · p. 13
  • Evaluation metrics include accuracy, AUROC, NLL, MCRMSE, MSE, RMLSE, MAE, RMSE, and MedAE.Source: §B.1 · p. 13
  • For each task, the authors provide a natural-language description and a Python random-guess baseline as the agents’ initial reference.Source: §B.1 · p. 14
  • The airline-reviews example is a regression task predicting British Airways review ratings from text, with ratings from 1 to 10 and RMSE evaluation.Source: §B.1 · p. 14
  • Its example pipeline uses a 10% train-validation split, evaluates validation predictions by RMSE, and submits predictions for test reviews.Source: §B.1 · p. 14

B.2. Details for Human Insight Case Collection · B.3. Model Configuration and Hyper-parameter Setting · C. Further Discussions

The supplied material explains how DS-Agent collected and structured Kaggle solution insights, gives examples of techniques in those cases, and specifies model and development-stage settings.

B.2. Details for Human Insight Case Collection

  • 12 recently completed Kaggle competitions—four each for text, time-series, and tabular data—supplied top-10 private-leaderboard team reports and top-10 public-leaderboard notebooks for the insight collection.Reports received basic text cleaning, while GPT-3.5 was prompted to extract textual solutions from code.Source: §B.2 · p. 15
  • The code-extraction prompt requested reproducible, code-grounded accounts of solution design, architecture, hyperparameters, objectives, advanced techniques, and performance tricks.It directed the model to focus on training and answer each question in detail.Source: §B.2 · p. 15
  • 0.453 in private on fold 3 came from merging similar essays and trying reverse-autocorrect augmentation in the showcased solution.The participant described this as a single-fold result and said they did not pursue further augmentation experiments.Source: §B.2 · p. 16
  • The example training workflow tuned hyperparameters on one fold, reused them on other folds, and ensembled the 10 best checkpoints across folds using inference maxlen 1500.The author chose 1500 for efficiency after trying values such as 1024 and 2048, then averaged checkpoints by fold and across folds.Source: §B.2 · p. 16
  • A showcased competition report said its author placed third and ranked DeBERTa, EMA, augmentation, differential learning rates, custom pooling, token_type_ids, and data cleaning as useful techniques.The author listed decoder models, AWP, FGM, weight decay, constant learning rate, handcrafted features, and GBT stacking as techniques that did not work.Source: §B.2 · p. 16

B.3. Model Configuration and Hyper-parameter Setting

  • The configurations used gpt-3.5-turbo-16k, gpt-4-0613, and Mixtral-8x7B-Instruct-v0.1; decoding temperature was 0.5 in development and 0.7 in deployment.Mixtral used vLLM for speedup (Kwon et al., 2023), and llm-embedder was the pretrained embedding model (Zhang et al., 2023b).Source: §B.3 · p. 17
  • Development-stage DS-Agent used 5 iterations, retrieved 5 cases, and set ndebug to 5; baseline reproduction followed ResearchAgent’s reported hyperparameters (Huang et al., 2023).Source: §B.3 · p. 17

C.1. Comparison with AutoML Techniques · C.2. Case Study

DS-Agent is compared with AutoGluon on tabular development tasks and is presented as more adaptable in its information sources, task support, and user interaction. The case study section introduces representative development and deployment examples.

C.1. Comparison with AutoML Techniques

  • In 2 of 4 tabular tasks, DS-Agent with GPT-4 outperformed AutoGluon by a large margin and performed on par with it in the other two.The comparison reports five DS-Agent runs, showing the best and average results; AutoGluon is reported for a single run. The GPT-3.5 average for spaceship-titanic is omitted because of a failed run.Source: Table 6 · p. 18, §C.1 · p. 18
  • DS-Agent can use updated public technical reports and Kaggle code to address data science tasks, whereas AutoML systems often require expertise, software development, and frequent updates.Source: §C.1 · pp. 17–18
  • DS-Agent dynamically builds and trains models for varied tasks, while AutoGluon does not naturally support enzyme-substrate multi-task classification without reframing it as multiple single-task problems.The passage contrasts DS-Agent’s flexibility with AutoML systems that typically constrain task types to tabular data.Source: §C.1 · p. 18
  • DS-Agent lets users describe tasks in natural language through a conversational interface, unlike traditional AutoML systems that require code and machine-learning expertise.The passage frames this interface as lowering the technical barrier to using advanced data science.Source: §C.1 · p. 18
  • The comparison covers four tabular development tasks because AutoGluon does not flexibly handle tasks such as time series classification and text regression.DS-Agent is run five times, with best and average results reported; AutoGluon is run once, and the GPT-3.5 spaceship-titanic average is omitted after a failed run.Source: §C.1 · p. 18, Table 6 · p. 18

C.2. Case Study

  • The case study presents two representative DS-Agent cases, covering development and deployment tasks.Source: §C.2 · p. 18

C.2.1. CASE STUDY IN DEVELOPMENT STAGE · C.2.2. CASE STUDY IN DEPLOYMENT STAGE · C.3. Error Mode Analyses

These sections examine DS-Agent’s development- and deployment-stage case studies, including classifier tuning and forecasting-code adaptation, and analyze traces and failure cases from both stages.

C.2.1. CASE STUDY IN DEVELOPMENT STAGE

  • 98.46% validation accuracy was the peak after switching to RoBERTa-large and recalibrating hyperparameters, following 90.59% after tuning learning rate and epochs.The earlier RoBERTa-base experiment reached 47.86% accuracy.Source: §C.2.1 · p. 18
  • Validation accuracy progressed from 75.21% with TF-IDF and logistic regression to 76.06% with BERT, before RoBERTa-base fell to 47.86%.The subsequent learning-rate and epoch tuning raised accuracy to 90.59%.Source: §C.2.1 · p. 18
  • Reported validation/test accuracies were 0.7521367521367521/0.9145299145299145 for the CountVectorizer regression model and 0.7606837606837606/0.7555555555555555 for BERT.Source: §C.2.1 · pp. 19, 21
  • Accuracy reached 0.9487179487179487 on the validation set and 0.9846153846153847 on the test set after switching to DeBERTa and tuning the learning rate from 1e-5.The experiment log identifies the model replacement and learning-rate tuning; the result reports the final validation and test accuracy.Source: §C.2.1 · p. 27

C.2.2. CASE STUDY IN DEPLOYMENT STAGE

  • In deployment, DS-Agent adapts the ETTm2 development-stage Bi-GRU/Bi-LSTM ensemble to the electricity forecasting task with minor Python-script modifications.The case study describes this reuse as reducing the foundational capabilities required of LLMs.Source: §C.2.2 · p. 27
  • For 96-step forecasts with seven input and output variables, the example trains bidirectional GRU and LSTM models and combines predictions using inverse validation-loss weights, evaluating MSE and MAE.The task description identifies MSE and MAE as evaluation metrics; the code trains both models and forms a validation-loss-weighted ensemble.Source: §C.2.2 · pp. 27, 29
  • For the Electricity dataset, with 96-step sequences and 321 variables, the example trains a bidirectional GRU and evaluates forecasts using MSE and MAE.The code submits predictions from that single model on the test set.Source: §C.2.2 · pp. 29, 31

C.3. Error Mode Analyses

  • The error-mode analysis examines DS-Agent traces and presents detailed failure modes for both development and deployment, alongside additional bad cases.The section points readers to a failure-case document in the project repository.Source: §C.3 · p. 31

C.3.1. ERROR MODE ANALYSES IN DEVELOPMENT STAGE

In development, DS-Agent with GPT-4 achieves a 100% success rate, while DS-Agent with GPT-3.5 may fail because of limitations in its foundational capabilities. GPT-3.5 failures involve unreasonable or overly complex plans and unsuccessful debugging, and stronger LLMs or further finetuning may mitigate the issue.

C.3.1. ERROR MODE ANALYSES IN DEVELOPMENT STAGE

  • 100% is the development-stage success rate of DS-Agent with GPT-4, while DS-Agent with GPT-3.5 may fail due to weaker foundational capabilities.The passage attributes GPT-4's success to its strong foundational capabilities.Source: §C.3.1 · p. 31
  • GPT-3.5 failures include unreasonable plans, plans too complex for the Programmer and Debugger, and debugging failures the Debugger cannot rectify.These are the three failure modes identified for DS-Agent in the development stage.Source: §C.3.1 · p. 31
  • Stronger LLMs or further finetuning for better alignment may mitigate DS-Agent with GPT-3.5's development-stage failures.The passage links the failures to GPT-3.5's relatively weaker foundational capabilities.Source: §C.3.1 · p. 32

C.3.2. ERROR MODE ANALYSES IN DEPLOYMENT STAGE · C.4. Some Discussions on DS-Agent · D. Prompt for DS-Agent

The deployment analysis identifies common code-generation failures, while the discussion addresses risks from reusing LLM-generated solutions and dependence on external resources. The prompt appendix specifies how DS-Agent’s components retrieve, plan, write, debug, log, and adapt code.

C.3.2. ERROR MODE ANALYSES IN DEPLOYMENT STAGE

  • Deployment failures include shape mismatches, undefined variables, incorrect function calls, and missing package imports across DS-Agent with GPT-4, GPT-3.5, and Mixtral-8x7b-Instruct.Source: §C.3.2 · p. 32
  • Other deployment failures include nonexistent dataframe keys, incompatible tensor data types, and incomplete programs; the authors suggest stronger LLMs or added debugging to address single-trial bugs.Source: §C.3.2 · p. 32

C.4. Some Discussions on DS-Agent

  • DS-Agent’s case bank may amplify LLM biases, but high-performance filtering, feedback-driven exploration, and adding human solutions are presented as ways to counter this risk.The authors argue that human solutions can contribute more innovative techniques and stronger empirical results than LLM-generated solutions alone.Source: §C.4 · p. 32
  • When external resources are unavailable, DS-Agent can rely on LLM knowledge or human-provided plans, but these options depend on model capability or human experts.The first option corresponds to the DS-Agent w/o CBR ablation; the second delegates experiment planning to humans and implementation to DS-Agent.Source: §C.4 · p. 32

D.1. RankReviser; D.2. Planner

  • RankReviser ranks five cases by relevance and helpfulness, while Planner uses a retrieved human-written case and experiment context to propose one incremental next step under stated constraints.Planner instructions include avoiding repeated techniques, multiple-trial plans, visualization, changes to data splits, and tuning more than two hyperparameters at once.Source: §D.1 · p. 33, §D.2 · p. 34

D.3. Programmer; D.4. Debugger

  • Programmer edits the provided script according to the plan and returns complete runnable code, while Debugger revises buggy code using the execution log and returns corrected code.The prompts specify execution with a 24 GB NVIDIA GPU and instruct both components to avoid unrelated changes.Source: §D.3 · p. 35, §D.4 · p. 35

D.5. Logger; D.6. Adapter

  • Logger summarizes the experiment and its observed result from instructions, execution output, and code changes, while Adapter generates complete code for a new task using example cases.Logger is instructed not to report guessed results; Adapter receives example tasks and code alongside the current task and script.Source: §D.5 · p. 36, §D.6 · p. 36
Loading 2402.17453v5…