🤖 AI

DoctorAgents: an agentic framework to iteratively refine AutoML pipeline for small clinical temporal data

DoctorAgents is an agentic AI framework that leverages specialized large language models to iteratively refine clinical machine learning pipelines through reasoning-driven updates and natural-language feedback, outperforming traditional AutoML systems on small, heterogeneous temporal data.

Original authors: Ruilin Wang, Bo-Hong Wang, Elizabeth Kourbatski, Jun Bai, Hegang Chen, Ziyang Song, Gilles Boire, Marie Hudson, Yue Li

Published 2026-08-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruilin Wang, Bo-Hong Wang, Elizabeth Kourbatski, Jun Bai, Hegang Chen, Ziyang Song, Gilles Boire, Marie Hudson, Yue Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to be a doctor. You give it a stack of patient records, hoping it can spot patterns that predict who might get sick or who needs a specific treatment. This is the world of Clinical Machine Learning. But here's the catch: real patient records are messy. They aren't neat spreadsheets where every row is a perfect snapshot in time. Instead, they are like a chaotic diary where a patient might visit the doctor once a year, then three times a week, then disappear for months. Some notes are missing, some are vague, and the timing of a measurement often matters just as much as the number itself.

To make sense of this chaos, scientists usually build Machine Learning Pipelines. Think of a pipeline as a custom-built factory assembly line. One part cleans the data, another part organizes the messy notes into a format the computer understands, and a final part builds the prediction model. Traditionally, building these factories has been a slow, manual job requiring a team of experts: a statistician to check the math, a computer scientist to write the code, and a doctor to make sure the logic makes medical sense. Recently, Artificial Intelligence (AI) has tried to automate this. Some systems, called AutoML, act like a frantic chef trying every possible recipe in the world to see which one tastes best. They throw thousands of combinations at the problem, hoping one sticks. But for small, messy medical datasets, this "brute force" approach often wastes time, misses the subtle clues, and produces results that are hard to understand.

Enter DoctorAgents, a new approach described in a paper by researchers at McGill University and the Quebec AI Institute. Instead of a frantic chef throwing ingredients at a wall, the researchers propose a team of specialized AI "doctors" who work together like a collaborative medical board. This system, called DoctorAgents, doesn't just guess; it reasons. It uses a team of Large Language Model (LLM) agents—specialized AI assistants—to analyze the data, build a pipeline, test it, and then think about why it failed before making a tiny, targeted fix.

The core idea is Iterative Agentic Refinement. Imagine you are solving a complex puzzle. A traditional computer might try to swap out the entire puzzle board every time a piece doesn't fit. DoctorAgents, however, acts like a detective. If a piece doesn't fit, it asks, "Why?" and then gently nudges just that one piece or changes the angle of the board slightly, rather than starting over. It uses a technique called Textual Gradient Descent, which is like translating a doctor's natural language feedback ("This feature is too noisy") into a precise code update without rewriting the whole program.

The researchers tested this system on four different clinical challenges using real patient data. These included predicting who would die in the ICU, who would be readmitted within a week, how long a patient would stay in the hospital, and whether a rheumatoid arthritis patient would respond to a specific drug treatment. The data was "small" in the grand scheme of AI (often just a few hundred to a few thousand patients) and "temporal," meaning the timing of every measurement was crucial and irregular.

The results suggest that DoctorAgents is a strong contender. In these experiments, the AI team consistently outperformed other automated systems, including the current state-of-the-art "brute force" methods and other AI agents. For instance, in predicting ICU mortality, DoctorAgents-DS (a specialized version of the system) achieved a score of 0.520, beating the next best method which scored 0.476. In predicting rheumatoid arthritis treatment response, it reached 0.577, surpassing the previous best of 0.564.

But the paper suggests something even more important than just a higher score: interpretability. Because the AI agents reason through their steps and keep a "memory log" of what worked and what didn't, the features they create make sense to humans. For the arthritis prediction, the system didn't just find a random number; it constructed features that a human doctor would recognize as meaningful, such as the variability in a patient's fatigue levels over a year or the specific type of drug they were prescribed. The system explicitly ruled out the idea that simply regenerating entire pipelines from scratch is the best way to go; instead, it found that targeted, memory-aware refinement leads to more stable and reliable results.

The authors are careful to note that while the system is robust and efficient, it is not a magic wand. In the rheumatoid arthritis task, the performance was modest across all methods, suggesting that the data itself might lack strong predictive signals, a limitation the AI cannot magically overcome. However, the study strongly suggests that for small, messy, time-based medical data, a team of reasoning AI agents that learn from their mistakes is a more effective strategy than a system that just tries everything and hopes for the best. It turns the process of building a medical prediction model from a game of chance into a thoughtful, iterative conversation between a machine and the data it is trying to understand.

Technical Summary: DoctorAgents

Problem Statement

Clinical machine learning (ML) faces significant hurdles when applied to small, heterogeneous, and temporally complex datasets, such as those found in rare-disease registries or single-center studies. Unlike standard tabular benchmarks, clinical records are often asynchronous, incomplete, and irregularly sampled, reflecting both underlying disease processes and the clinical observation process. Traditional automated machine learning (AutoML) systems, which rely on brute-force search over predefined spaces (e.g., Bayesian optimization, evolutionary strategies), struggle in this domain. They often require substantial computational resources, lack explicit reasoning capabilities regarding data properties (such as leakage risks or temporal validity), and fail to provide interpretable, task-specific patient representations. Furthermore, existing Large Language Model (LLM)-based AutoML agents often operate in a trial-and-error regime with limited persistent memory, leading to unstable optimization, repetitive errors, and inefficient compute usage when refining pipelines for small clinical datasets.

Methodology: DoctorAgents

The authors propose DoctorAgents, an agentic AI framework that reformulates AutoML for small clinical temporal data from exhaustive search to reasoning-driven refinement. The system employs a multi-agent architecture where specialized LLM agents collaborate to construct, validate, and iteratively optimize end-to-end ML pipelines.

Core Agent Architecture

The framework utilizes a team of specialized agents, each with a distinct role:

Consultant ( $A_c$ ): Acts as a high-level analytical expert, transforming raw dataset metadata and user inputs into a structured task specification with domain knowledge to guide downstream synthesis.
Generator ( $A_g$ ): Translates the task specification into an initial executable ML pipeline code, covering data loading, preprocessing, model selection, training, and evaluation.
Repairer ( $A_r$ ): Handles execution failures by minimally modifying code to restore executability without altering the original algorithmic intent.
Logger ( $A_l$ ): Maintains a structured, long-term memory log of inter-iteration changes, performance outcomes, and meaningful differences between pipeline versions. It employs a summarization strategy (SUMMA) to manage context length and prevent "note explosion."
Evaluator ( $A_e$ ): Assesses the pipeline's quality, correctness, and potential flaws (e.g., overfitting, leakage) using natural language feedback based on performance metrics and historical logs.
Optimizer ( $A_o$ ): Converts the evaluator's natural language feedback into concrete code modifications using Textual Gradient Descent (TGD). This allows for targeted, localized updates to the pipeline without regenerating the entire codebase.

Workflow

The system operates in a closed-loop iterative process:

Initialization: The Consultant analyzes metadata to define the task, and the Generator synthesizes an initial pipeline.
Execution & Verification: The pipeline is executed. If it fails, the Repairer fixes it.
Memory Composition: The Logger records the changes and outcomes, summarizing history at checkpoints to maintain context efficiency.
Refinement: The Evaluator critiques the pipeline, and the Optimizer applies TGD to update the code based on this feedback.
Iteration: This cycle repeats until a stopping criterion is met.

Domain-Specific Specialization (DoctorAgents-DS)

To address the unique challenges of irregular temporal clinical data, the authors introduce a specialized variant, DoctorAgents-DS. This variant decomposes the optimization into two dedicated subtasks managed by specialized agents:

Data Preprocessing (DP) Agents: Focus exclusively on feature engineering, temporal aggregation, and handling missingness for irregular time-series data.
Model Development (MD) Agents: Focus on model selection, hyperparameter tuning, and training strategies.
This separation allows for deeper exploration of task-specific temporal representations while preserving explicit, editable code.

Key Contributions

The paper claims the following contributions:

Reasoning-Driven Agentic AutoML: A framework specifically designed for small clinical temporal datasets that combines dataset reasoning, execution validation, persistent memory, and iterative refinement.
Memory-Aware Optimization: A mechanism enabling long-term learning across iterations to prevent repeated failures and ensure stable refinement.
Domain-Specialized Design: A decomposition strategy (DP and MD agents) tailored for irregularly observed temporal clinical data, enabling richer, task-specific patient representations.
Automatic Clinical Feature Composition: The ability to automatically compose clinically meaningful variables from raw temporal data, moving beyond fixed preprocessing templates.
Empirical Validation: Extensive experiments demonstrating robustness, efficiency, and interpretability superior to existing methods.

Experimental Results

The authors evaluated DoctorAgents on four temporal clinical prediction tasks:

Datasets: Three tasks from the public MIMIC-IV database (ICU mortality, ICU readmission, and ICU length-of-stay prediction) and one task from an in-house Early Undifferentiated PolyArthritis (EUPA) cohort (predicting anti-TNF treatment response in rheumatoid arthritis).
Baselines: Comparisons included traditional ML (TabPFN), general-purpose LLMs (GPT-5, GPT-oss, Llama-8B), and existing agentic AutoML systems (AutoML-Agent, ERA).
Performance:
- DoctorAgents-DS achieved the strongest mean performance across all benchmarks, outperforming all baselines in terms of AUROC and AUPRC.
- The standard DoctorAgents framework also consistently outperformed established AutoML and agentic baselines.
- Ablation studies confirmed that both the Consultant (improving initialization) and Memory (stabilizing refinement) modules are critical for performance, with the full system achieving significantly higher validation scores than variants lacking these components.
Interpretability:
- The system generated features that were more technically valid and semantically consistent with clinical contexts compared to baselines.
- An independent evaluation using GPT-5.4 as a judge showed that DoctorAgents-DS produced a significantly higher proportion of "Well Justified" and "Clinically Reasonable" features (mean improvements of ~7-8% over baselines).
- The generated features aligned with clinical intuition (e.g., capturing physiological instability, monitoring intensity, and treatment trajectories) rather than relying on generic aggregations.

Significance and Claims

The paper posits that DoctorAgents represents a paradigm shift in clinical ML automation. By replacing brute-force search with a stateful, reasoning-driven process, the framework addresses the specific constraints of small, irregular clinical datasets. The authors claim that the system not only improves predictive performance but also produces interpretable, task-specific representations grounded in feature semantics and clinical rigor.

The significance lies in the framework's ability to autonomously handle the complexities of temporal clinical data—such as asynchronous observations, missingness patterns, and the distinction between physiological change and monitoring artifacts—without requiring manual, domain-expert pipeline construction. The results suggest that agentic AI, when equipped with persistent memory and domain-specialized reasoning, can serve as a robust tool for developing reliable clinical prediction models, potentially reducing the burden on data scientists and clinicians while improving the quality and interpretability of the resulting models.