DSMentor: Enhancing Data Science Agents with Curriculum Learning and Online Knowledge Accumulation

TL;DR

DSMentor introduces a curriculum learning framework for LLM data science agents, organizing tasks by difficulty and utilizing long-term memory to improve performance and reasoning on complex problems.

Contribution

This work presents a novel inference-time optimization method that applies curriculum learning and knowledge accumulation to enhance LLM agent effectiveness in data science tasks.

Findings

Up to 5.2% improvement in pass rate on benchmarks.

8.8% better causality reasoning compared to GPT-4.

Effective knowledge retention enhances problem-solving ability.

Abstract

Large language model (LLM) agents have shown promising performance in generating code for solving complex data science problems. Recent studies primarily focus on enhancing in-context learning through improved search, sampling, and planning techniques, while overlooking the importance of the order in which problems are tackled during inference. In this work, we develop a novel inference-time optimization framework, referred to as DSMentor, which leverages curriculum learning -- a strategy that introduces simpler task first and progressively moves to more complex ones as the learner improves -- to enhance LLM agent performance in challenging data science tasks. Our mentor-guided framework organizes data science tasks in order of increasing difficulty and incorporates a growing long-term memory to retain prior experiences, guiding the agent's learning progression and enabling more effective utilization of accumulated knowledge. We evaluate DSMentor through extensive experiments on DSEval and QRData benchmarks. Experiments show that DSMentor using Claude-3.5-Sonnet improves the pass rate by up to 5.2% on DSEval and QRData compared to baseline agents. Furthermore, DSMentor demonstrates stronger causal reasoning ability, improving the pass rate by 8.8% on the causality problems compared to GPT-4 using Program-of-Thoughts prompts. Our work underscores the importance of developing effective strategies for accumulating and utilizing knowledge during inference, mirroring the human learning process and opening new avenues for improving LLM performance through curriculum-based inference optimization.