🤖 machine learning

DrugSAGE:Self-evolving Agent Experience for Efficient State-of-the-Art Drug Discovery

DrugSAGE is a self-evolving agent framework that accumulates and reuses cross-task experiences—such as verified skills, effective strategies, and error fixes—to efficiently build state-of-the-art drug discovery models, significantly outperforming existing agents by reducing or eliminating the need for costly search on new tasks.

Original authors: Yikun Zhang, Xiwei Cheng, Tianyu Liu, Yuanqi Du, Wengong Jin

Published 2026-05-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yikun Zhang, Xiwei Cheng, Tianyu Liu, Yuanqi Du, Wengong Jin

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 bake the perfect cake for a specific, difficult recipe. In the world of drug discovery, this "cake" is a computer model that predicts how a new medicine will behave in the human body.

Currently, if a scientist wants to bake a cake for a new recipe (a new drug property), they usually start from scratch. They have to guess which ingredients to use, try different oven temperatures, taste the batter, fail, and start over. Even if they bake 100 cakes for 100 different recipes, they often throw away all the lessons learned after the first one. They treat every new recipe as a brand-new mystery, wasting time and money re-learning things they already know.

DRUGSAGE is a new "smart baker" (an AI agent) that changes the rules. Instead of forgetting its mistakes and successes, it keeps a giant, evolving recipe book that gets smarter with every cake it bakes.

Here is how it works, broken down into simple parts:

1. The "Recipe Book" (The Memory)

Most AI agents have a short-term memory; they forget what they did yesterday. DRUGSAGE has a long-term memory made of three specific sections:

The "Winning Recipes" (Solution Memory): It remembers which specific combinations of ingredients (models and training methods) worked best for similar tasks in the past.
The "Chef's Notes" (Refinement Memory): It remembers why a change worked. For example, "When we lowered the oven temperature by 10 degrees, the cake rose better."
The "Disaster Log" (Execution Memory): It keeps a list of things that caused the kitchen to catch fire (code errors, memory crashes) and exactly how to fix them next time.

2. The "Smart Search" (How it finds the best model)

When DRUGSAGE is asked to bake a new cake, it doesn't just guess randomly. It uses a strategy called MCTS (Monte Carlo Tree Search), which is like exploring a forest of possible recipes.

Without Memory: An agent wanders the whole forest, trying every path, even the ones that lead to dead ends.
With DRUGSAGE: Before it even steps into the forest, it checks its Recipe Book. It says, "Hey, for a cake similar to this one, the 'chocolate-chip' path worked great last time." It skips the dead ends and goes straight to the promising paths.
The Guarantee: The paper mathematically proves that even with this "cheat sheet" from the past, the agent is still guaranteed to find the absolute best solution eventually, just much faster.

3. The "Zero-Search" Superpower

This is the most impressive trick. Sometimes, the new recipe is so similar to an old one that DRUGSAGE doesn't need to bake at all.

The Scenario: If you ask it to predict how a drug dissolves in water, and it has already solved a very similar problem about how a drug dissolves in oil, it might say: "I already know the answer. I'll just copy the winning recipe from the oil task and use it here."
The Result: It delivers a top-tier solution with zero extra testing or searching. It's like a chef who, when asked to make a lemon tart, realizes they already made a perfect lime tart yesterday and just swaps the fruit, saving hours of work.

The Results: Did it work?

The authors tested this "smart baker" on 33 different drug discovery challenges (like predicting toxicity, solubility, or how well a drug binds to a protein).

Solo Performance: When forced to work alone without help, DRUGSAGE still beat all other top AI agents.
Teamwork Performance: When allowed to use its memory from 16 smaller tasks to help with 17 new, harder tasks, it crushed the competition.
The Big Win: In the "Zero-Search" mode (where it just copies past solutions without testing), it outperformed other agents that were allowed to spend hours searching for answers. It was 10% to 30% better than the competition, even though the competition was working much harder.

In a Nutshell

DRUGSAGE is an AI that learns to learn. Instead of treating every drug discovery problem as a fresh start, it treats every problem as a lesson that makes the next one easier. It builds a library of "verified skills" and "proven fixes," allowing it to solve complex drug problems faster, cheaper, and more accurately than any previous system.

Technical Summary: DRUGSAGE

Problem Statement

Building state-of-the-art (SOTA) predictive models for drug discovery is an expensive search problem. Current Large Language Model (LLM)-based agents can achieve SOTA performance by extensively searching over tools, architectures, and training strategies. However, these agents typically treat each task as an isolated problem, discarding all accumulated search experience (e.g., effective featurizations, training recipes, or error fixes) once a task is completed. Consequently, they must pay the full computational and token search cost for every new task, making it difficult to scale SOTA model development across the thousands of diverse prediction problems found in drug discovery.

Methodology

The authors propose DRUGSAGE (Self-evolving AGent Experience), a framework designed to accumulate and reuse experience across tasks to build SOTA drug discovery models efficiently. The system operates through three primary components:

1. Executable Skill Library ( $K$ )

Before budgeted search begins, an Explore Agent constructs a shared skill library by searching literature and GitHub repositories. This process involves:

Discovery: Querying literature from multiple expert perspectives to identify relevant papers and repositories.
Grounding: Cloning repositories and extracting Abstract Syntax Tree (AST) snapshots to ensure the LLM only uses symbols observed in the code.
Validation: Running tiered checks (syntax, import resolution, end-to-end execution in sandboxes) with automatic repair. Only validated skills are added to $K$ .

2. Cross-Task Experience Memory ( $Z$ )

DRUGSAGE maintains a persistent memory $Z = (R, H, Q)$ that stores evidence from previous tasks:

Solution Memory ( $R$ ): Records the performance of model families and specific solutions across different tasks, including architecture, hyperparameters, and validation scores.
Refinement Memory ( $H$ ): Stores information about edges in the search tree, including LLM-proposed edits, rationales, and the performance deltas between parent and child nodes.
Execution Memory ( $Q$ ): Logs execution failures, verified fixes, resource profiles (runtime/memory), and environment traces to enable quick repairs during future runs.

3. Experience-Conditioned Search

The core search loop utilizes a Memory-Enhanced Monte Carlo Tree Search (MCTS) when a search budget ( $B > 0$ ) is available. The memory $Z$ guides four stages:

Root Selection: Uses a memory-augmented Upper Confidence Bound (UCB) rule to select model families, balancing current task performance with historical transfer scores. The authors prove this modification preserves the $O(|F| \log T)$ regret bound of standard UCB.
Parent Sampling: Selects specific parent solutions to expand based on target task performance, expansion breadth, and historical performance.
Child Generation: The LLM proposes edits grounded by $H$ , retrieving relevant rationales and successful recipes from similar past tasks.
Execution & Repair: Uses $Q$ to prepare environments and apply verified fixes immediately upon encountering known errors.

4. Zero-Test-Time Search (DRUGSAGE-ZERO)

When the search budget is zero ( $B=0$ ), the system bypasses search entirely. It forms a task signature (task type, metric, dataset size, description embedding) to identify analogous tasks in the memory $Z$ . It retrieves the best verified solution from the most similar historical task and deploys it directly on the new target.

Key Contributions

Framework Design: Introduction of DRUGSAGE, an autonomous agent that integrates cross-task memory into an MCTS-based search loop. The authors provide a formal guarantee that the memory-augmented selection policy preserves the regret bound of standard UCB.
Automated Skill Discovery: Unlike prior works (e.g., Autoresearch) that refine user-provided or top-leaderboard models, DRUGSAGE automatically builds a skill library by searching literature and repositories, removing the need for human-curated starting points and broadening the search space.
Zero-Test-Time Efficiency: Demonstration that cross-task memory enables a "zero-test-time search" regime. DRUGSAGE-ZERO transfers verified solutions to new tasks without any test-time search, outperforming baseline agents that perform 20 search iterations.

Experimental Results

The framework was evaluated on 33 molecular property prediction tasks across two benchmarks: Therapeutic Data Commons (TDC) and Polaris Hub.

Single-Task Setting: In a setting where agents build solutions from scratch without cross-task memory, DRUGSAGE ranked first among nine SOTA agents (including Autoresearch, STELLA, and Biomni) with an average rank of 1.95 and 10 wins out of 22 tasks.
Cross-Task Setting: Using memory accumulated from 16 smaller tasks, DRUGSAGE was evaluated on 17 held-out tasks.
- DRUGSAGE-ZERO ( $B=0$ ): Achieved an average normalized score of 0.935 on 17 held-out tasks, outperforming all baseline agents (which utilized up to 20 search steps) by 10–30%.
- Efficiency: DRUGSAGE-ZERO incurred near-zero LLM API costs at test time, whereas baselines required substantial compute and token budgets for iterative search.
- Performance: In specific cases (e.g., Solubility and CYP2D6 Inhibition), the zero-shot transfer solution outperformed the full search version of DRUGSAGE, suggesting that the transferred strategy was superior to what the agent could discover from scratch within the budget.

Significance and Claims

The paper claims that DRUGSAGE demonstrates the distinct advantage of cross-task memory for efficient SOTA model development in drug discovery. By treating each task as an experience that enriches the agent for future tasks, the system narrows the search space and can, in some cases, bypass search entirely. The work highlights that many drug discovery tasks share structural similarities (input/output structure, featurization needs) that current agents overlook, and that accumulating verified skills and statistical evidence allows for rapid adaptation to new problems without incurring the full cost of iterative trial-and-error. The authors note that while the framework is tested primarily on drug discovery, the approach of accumulating structured cross-task experience could be generalized to other scientific research areas.