DrugMCTS: a drug repurposing framework combining multi-agent, RAG and Monte Carlo Tree Search

In plain words

Drug repositioning asks which existing protein targets a known drug molecule might act on, so that the drug could be considered for a new use. A general language model asked this cold is unreliable, and training a specialised model for every drug area is expensive. DrugMCTS keeps a general language model as it is and gives it helpers. Similar molecules are looked up in a database, their chemical properties are computed, and notes on the known protein pockets and abstracts about the candidate targets are gathered and summarised for the model.

A search procedure then tries many different chains of these steps. Each finished chain gets a score from two checks: whether the model gives the same answer when asked several times, and whether it says yes when asked directly whether the molecule interacts with the chosen target. The search keeps returning to chains that scored well, but also tries options it has visited less, so it does not settle on the first decent answer.

The design lets the model reason over text summaries of structured data rather than raw molecular structures, and lets the search repeat and refine its decision. The paper reports that this small model, with no training for the task, beats larger general models and trained deep learning baselines on its benchmark.

TL;DR

• DrugMCTS uses Qwen2.5-7B-Instruct without fine-tuning, inside Monte Carlo tree search (MCTS) with five agents that together perform six actions §2.3 alg. 1.
• Each rollout's reward is the mean of a relative score (how often the model picks the same target) and an absolute score (share of yes answers) eq. 4 eq. 5 eq. 6.
• Author-reported recall on DrugBank and KIBA is 0.4466 and 0.4224 with the "GT count" selection, and 0.5534 and 0.4924 with "GT + 3" tab. 1. Both settings use the ground-truth count, so check the protocol before comparing.
• The official repository is a separate MIT-licensed tool. Its own README says it does not reproduce the paper's numbers and differs from the experiments README.md README.md.

Problem

Fine-tuning a large model on domain data is costly, must be repeated as the science changes, and risks catastrophic forgetting §1. Retrieval-augmented generation avoids that cost, but in drug-target work it usually retrieves text or knowledge-graph facts and leaves out molecular structure and protein data §1. General-purpose models also drop information from the middle of long inputs ("lost in the middle") and typically answer in one shot, with no feedback loop §1 fig. 1. The authors state that deep learning models can lose more than 20% accuracy on molecule-protein pairs unlike their training data §2.4 §1.

Method

Overview. Given a query molecule, DrugMCTS coordinates a sequence of agents that retrieve similar molecules, analyse them, filter them, analyse the protein pockets and pick a target. MCTS guides the search over these sequences §2.1 fig. 3 alg. 1.

MCTS. Each search step runs four phases: selection by UCT, expansion, simulation, and backpropagation. Each rollout reaches an end node, is scored, and updates every node on its path §2.2 alg. 2. After a set number of rollouts, the paper says the final answer is chosen from the candidates, but does not say how §2.2.

# MCTS for drug repositioning  [alg. 2] [§2.2]
# Input:  query molecule, rollout budget N (12 in the experiments [§3.2])
# Output: selected target(s) with their rewards

def drugmcts(query, rollouts):
    # The root holds the query molecule only  [alg. 2]
    root = Node(state={"query": query})
    for _ in range(rollouts):
        # Selection: descend by UCT to a leaf  [§2.2] [eq. 1]
        leaf = select_by_uct(root)
        # Expansion and simulation: run the agent actions to an end node  [§2.2]
        end_node = rollout(leaf)
        # Backpropagation: update visits and rewards on the path  [§2.2]
        backpropagate(end_node, end_node.reward)
    # Algorithm 2 returns one best child by average reward  [alg. 2]
    return best_child(root)

# UCT score of child i  [eq. 1]
# W_i: summed reward of the child, n_i: its visits, N: parent visits,
# c: exploration constant (value not stated in the paper)

def uct_score(child, parent_visits, c):
    # Exploitation: mean reward so far  [eq. 1]
    exploitation = child.total_reward / child.visits
    # Exploration: favours rarely visited children  [eq. 1]
    exploration = c * sqrt(log(parent_visits) / child.visits)
    return exploitation + exploration

Action space. The paper defines six actions. Five agents (retrieval, molecule analysis, molecule selection, interaction analysis, decision) perform them. The retrieval agent calls no model; the others use the same LLM §2.3. A6 is the end action, which calls no agent §2.3.

A1 Retrieval. Compute the Tanimoto similarity and a ChemBERTa cosine similarity between the query and every molecule in the database. Take the top 10 under each metric, merge, and remove duplicates, giving the candidate molecules. Proteins known to interact with any candidate form the candidate protein set §2.3 eq. 2 eq. 3. Because the scores are deterministic, the pool is the same on every run, which the paper needs for a unique root §2.3.

# One rollout through the actions  [alg. 1] [§2.3]
# Input:  query molecule (SMILES)
# Output: selected target P_s and its final reward

def drugmcts_rollout(query):
    # A1: top 10 by Tanimoto and top 10 by ChemBERTa cosine, merged and
    # de-duplicated  [§2.3] [eq. 2] [eq. 3]
    candidates = dedup(top10_tanimoto(query) + top10_cosine(query))
    # Proteins known to interact with any candidate  [§2.3] [alg. 1]
    candidate_proteins = proteins_of(candidates)

    # A2: properties of the query and a text report  [§2.3] [alg. 1]
    query_props = molecular_properties(query)
    query_report = llm_summarize(query, query_props)

    # A3: the model keeps reference molecules; candidates get properties
    # only, without a report  [§2.3] [alg. 1]
    references = llm_select(query, query_report, candidates)
    # Reference proteins: candidate proteins that bind a reference  [§2.3]
    reference_proteins = proteins_of(references) & candidate_proteins

    # A4: pocket text from PLIP and abstracts for reference proteins  [§2.3]
    pockets = plip_pocket_text(reference_proteins)
    literature = pubmed_abstracts(reference_proteins)
    interaction_report = llm_summarize_interactions(
        references, reference_proteins, pockets, literature
    )

    # A5: decision agent picks one target  [§2.3] [alg. 1]
    selected = llm_select_protein(
        query, query_report, references, reference_proteins,
        pockets, interaction_report,
    )
    # A6: end action; the rollout stops here and is scored  [§2.3] [alg. 1]
    return reward_of(selected, query, query_report, pockets, literature)

A2 Molecule analysis. Compute structural and physicochemical properties of the query with RDKit and PubChemPy. The listed properties are chiral centres, scaffolds, functional groups, molecular weight, logP, polar surface area, hydrogen-bond donors and acceptors, rotatable bonds and heavy-atom count. The agent turns them into a report §2.3. The same properties are computed for candidates, without a report, as part of A3 §2.3.

A3 Molecule selection. The model filters the candidate pool by structural similarity, pharmacophore integrity and drug-likeness, which the paper says removes noise and helps against "lost in the middle" §2.3 §1. The reference proteins are the candidate proteins that bind a selected reference §2.3.

A4 Interaction analysis. Pocket data come from PLIP applied to PDB structures. The tabular output is rewritten as descriptive sentences. PubMed literature supplies context §2.3 alg. 1. The authors adopt this from DrugRealign §2.3.

A5 Protein selection. The decision agent reads everything above and selects one target. The text says it chooses "from the full list of candidates" §2.3; Algorithm 1 passes it the reference proteins alg. 1. See Pitfalls.

Reward. The paper uses two signals. Each rollout's selection is repeated k times, and the most frequent answer p* is the relative reward. A decision model is then asked k times whether the chosen protein interacts with the query. The share of yes answers is the absolute reward. The final reward is their average §2.4 alg. 1 eq. 4 eq. 5 eq. 6. The self-consistency idea comes from a cited work §2.4.1.

# Final reward of the chosen target  [§2.4] [alg. 1]
# k = 4 samples per rollout  [§3.2]

def reward_of(selected, query, query_report, pockets, literature):
    # Re-run the selection k times; p* is the most frequent answer  [§2.4]
    picks = [llm_select_protein(...) for _ in range(k)]
    best = most_common(picks)
    # Relative reward: share of picks equal to p*  [eq. 4]
    relative = count(picks, best) / k
    # Ask k times whether the query interacts with the target  [§2.4]
    judgements = [llm_judge_interaction(best, ...) for _ in range(k)]
    # Absolute reward: share of yes answers  [eq. 5]
    absolute = count(judgements, "yes") / k
    # Equal-weight average of the two signals  [eq. 6]
    return (relative + absolute) / 2

The ... marks inputs that the paper lists in §2.4 and alg. 1. Which target the absolute judgement sees (the single A5 output or p*) is not explicit in the paper; see Pitfalls.

Evaluation metric. Recall: correctly predicted proteins divided by all ground-truth proteins for an entry §3.1 eq. 7.

Hyperparameters

No model is trained by this method. The table lists the values the paper and its repository state, with the source of each.

Results

Author-reported. Setup: DrugBank and KIBA, with 788 and 626 entries §3.1 tab. 3; Qwen2.5-7B-Instruct with 12 rollouts §3.2; baselines are general models with and without RAG, and four deep learning models §3.3. Hardware is not stated.

Headline comparison tab. 1:

The "Selection" column is the paper's label. Both DrugMCTS rows are the protocol the paper uses for its headline numbers; the paper does not define "GT count" at prediction time tab. 1 §3.4.

Ablations on the same benchmarks, with Qwen2.5-7B-Instruct, each cell "top-k / top-k+3" where the paper reports both tab. 2:

Other reported results:

• Rollouts: going from 8 to 12 gives clear gains on both datasets; going from 12 to 24 helps only the KIBA top-k+3 case. The paper keeps 12 §4.3 fig. 5 fig. 6. The figures are plots, and their values are not in the text.
• Case study: Equol with CXCR3 (PDB 8K2W) has a predicted AutoDock Vina affinity of −8.4 kcal/mol. The paper calls it an interaction not reported before, supported by docking and PyMOL visualisation; no wet-lab test is described app. A.2 fig. 7 fig. 8.
• Costs: token counts appear only in a figure fig. 6; the paper gives no per-query money or time figure.

Paper vs code

The repository is https://github.com/yaoge777/DrugMCTS-tool, linked from the paper, and licensed MIT. Its README says it is a standalone tool that does not import or modify the experimental project's code, and that it "is not a claim that the paper's reported benchmark numbers were reproduced" README.md README.md. The experimental pipeline behind the paper's tables is not in this repository. Code citations below are to src/drugmcts/ at commit 23d714776a68f44425de0e10e82d1f01992c99c4.

• Retrieval default. The paper uses Tanimoto plus ChemBERTa §2.3. The default config uses Tanimoto only; the union is an opt-in "hybrid" mode config.py retrieval.py. The README says the default must not be called the paper's dual retrieval README.md.
• ChemBERTa vector. The paper says "last hidden state" §2.3. The code takes the first token of the last hidden state retrieval.py. Inputs are truncated at 128 tokens retrieval.py. The model is DeepChem/ChemBERTa-77M-MLM, which the paper does not name config.py.
• Fingerprint. Not stated in the paper. Code uses Morgan radius 2 with 2048 bits chemistry.py.
• Molecule properties. The paper lists functional groups and uses RDKit and PubChemPy §2.3. The code computes RDKit descriptors: molecular weight, logP, TPSA, donors, acceptors, rotatable bonds, heavy atoms, scaffold and chiral centres chemistry.py. It computes no functional-group list and makes no PubChem call; cached PubChem fields come from an import step engine.py legacy.py.
• Pockets and literature. The paper runs PLIP and queries PubMed §2.3. The code reads pre-computed PLIP outputs and cached abstracts from an import folder, with no live PLIP or PubMed calls legacy.py legacy.py. It keeps at most 5 entries per PLIP interaction type and 4000 characters per abstract. The paper states neither limit legacy.py legacy.py.
• Protein pool for selection. Algorithm 1 passes the reference proteins to the decision agent alg. 1. The code restricts the choice to proteins of the selected references engine.py.
• End action. The paper treats A5 and A6 as separate actions §2.3. The code creates the terminal node in the same step as the protein selection, and treats a node at stage 4 as terminal engine.py engine.py.
• Selection ties. The paper says p* is the most frequent answer §2.4 and does not break ties. The code breaks them by protein identifier (lexical order) engine.py.
• Which target is judged. The paper's Algorithm 1 feeds P_s to the absolute judgement and uses p* for the relative reward alg. 1. The code judges the majority pick p* engine.py.
• Invalid responses. The paper does not say what happens to malformed answers. The code drops them and retries, up to three attempts, and marks the path failed if fewer than k valid answers arrive engine.py.
• Repeated visits. A terminal already scored is reused without a new vote. Each revisit backpropagates the stored reward engine.py. The paper does not say this.
• Selection rule. The code expands unvisited children first, then applies UCT with parent visits floored at 1 engine.py. The paper gives the UCT formula only eq. 1 and does not mention unvisited-first.
• Simulation. The paper names a simulation step but does not define it §2.2. The code finishes each rollout by expanding each new node and picking a random child until a terminal node engine.py.
• Ranking and output. Algorithm 2 returns one best child alg. 2. The code returns up to top_k targets, ranked by the mean reward over unique terminal paths, with ties broken by path count and then identifier engine.py engine.py. The number of targets comes from the caller, with no ground-truth input mcp_server.py.
• Budgets and partial results. The code adds a 120-request and 900-second budget. When a budget runs out, it returns the paths it has, labelled "partial" config.py engine.py. The paper does not mention budgets.
• Ablations. The paper reports six ablation setups S1–S6 §4.1. The code has no switches for excluding the analyses, relative-only rewards or the baseline setup. The ablations cannot be re-run from this repository.
• Dataset construction. The paper's filters and split for the benchmark app. A.1 are not in the repository. The importer keeps only positive labels and proteins with PLIP caches legacy.py legacy.py.
• Protein identity. The importer uses PDB structure identifiers as target identifiers, and records that they are structural representatives, not unique proteins legacy.py legacy.py. The paper does not say how targets are identified.
• Prompt limit. Over-long prompts raise an error and are not truncated engine.py. The paper does not describe this.
• Single model. The code uses the same configured LLM for every action and the judgement engine.py. The paper says other agents use the same LLM §2.3 but calls the absolute-reward model a "decision-making model" §2.4.
• Reproduction. The code records paper_result_reproduction: False engine.py and the README says the same README.md.

No discrepancies were found in the reward formulas (eq. 4–6) or the UCT form (eq. 1) as implemented engine.py engine.py.

Pitfalls

• Protocol: the headline numbers depend on the ground-truth count. The two DrugMCTS rows of tab. 1 are labelled "selection = GT count" and "GT + 3". The paper does not define how many targets a prediction returns when the count is unknown. The repository's README says the tool "does not require the number of known targets" README.md. Report recall at a fixed top-k as well, and say which protocol you used.
• Text and table disagree. The text gives the deep-learning KIBA recall as 26.45% §3.4; Table 1 gives 0.3216 tab. 1. The text says the ensemble is "84% better" than the best general model on DrugBank §3.4; Table 1 gives 0.2364 / 0.1619 = 1.46, or 46% tab. 1. The text's "330%" and "91.4%" over general models §3.4 do not match one baseline in Table 1 (recomputed: 340% over 0.1259 on DrugBank; 91% over 0.2580 on KIBA). Use the table values.
• Retrieval embedding. §2.3 uses ChemBERTa for molecule cosine similarity (A1). §3.2 says the retrieval phase uses ESM2_t33_650M_UR50D, a protein language model, for cosine similarity. Both cannot be right for molecule retrieval. The code uses ChemBERTa config.py.
• Undefined symbol. Equation 3 uses sim_T, which is never defined as an equation; only the text describes it §2.3 eq. 3.
• Candidate pool wording. A5 says "the full list of candidates" §2.3; Algorithm 1 uses the reference proteins alg. 1. The code uses the second.
• Cross-references. The text refers to "Appendix Table 1" and "Appendix Table 2" for dataset statistics §3.4. In the compiled numbering the dataset statistics are Table 3 and the baseline ratios are Table 4 tab. 3 tab. 4. Limitations cite "Section 2.4" for pocket data §6; pockets are described under A4 in §2.3 §2.3. Use the tables and section headings as printed in this card.
• Ensemble names. The deep-learning baselines are named AttentionDTA and GraphDTA in §3.3, and AttentionDTI and GINConvNet in §3.4.
• Unverifiable count. The conclusion cites "1,221 experimental interactions" §5. No table in the paper gives that number. Tables 3 to 5 give other counts.
• Reward granularity. With k = 4, each of the relative and absolute scores takes one of five values (0, 0.25, 0.5, 0.75, 1), so the final reward takes nine values in steps of 0.125 (derived from eq. 4–6). Ties are common, and the ranking tie-break decides many results.
• Exploration constant. The paper does not give c. With rewards between 0 and 1, the code's default c = 3.0 config.py makes the exploration term large relative to reward differences. This is an inference from the code's default, not something the paper says.
• Modest reward gain. The limitations say the combined reward gives about 1% over relative-only §6. Table 2 gives 1.46 points on DrugBank and 0.31 points on KIBA for the same comparison tab. 2. The gain is small and the benchmark is small (788 and 626 entries) §3.1.
• Randomness. The tree search is seeded (42 in the code), but the LLM samples at temperature 0.8 are not seeded, so reruns can differ config.py llm.py.
• Cost scales with the tree. One expansion is one chat request asking for up to four choices llm.py. A newly created path needs up to five expansions and reward calls, and failed responses trigger retries engine.py. The paper reports no timing or money figures.
• Retrieval depends on a stored database. The benchmark needs pocket and literature data prepared from PLIP and PubMed caches. The code's importer depends on those caches, and the repository ships only synthetic demo data README.md legacy.py.

Lineage

• Builds on: self-consistency sampling for the relative reward §2.4.1 [/papers/2203.11171]. Monte Carlo tree search (Chaslot, 2010) §2.2, UCT (Couëtoux et al., 2011) eq. 1 and DrugRealign (Wei et al., 2024) for the pocket-text method §2.3 have no card here.
• Compared against: GPT-4o-mini and DeepSeek-R1, with and without retrieval §3.3, and four deep-learning models (AttentionDTA, GraphDTA, DeepConv_DTI, Perceiver CPI) §3.3.
• Superseded by: not stated.