Probe Before You Edit: Probing-Guided Molecular Optimization for LLM Agents in Structure-Based Drug Design

The paper introduces PROBE, a probing-guided framework for LLM agents in structure-based drug design that decomposes ligands into editable sites to create a site map and EditManual, thereby enabling iterative multi-agent optimization that effectively balances binding affinity and druggability while mitigating the failure modes of existing methods.

The Gist

The Big Problem: Guessing in the Dark

Imagine you are a chef trying to fix a soup that tastes a bit off. You want it to be flavorful (like a drug binding tightly to a protein) but also healthy and easy to digest (like a drug being safe and easy to make).

Currently, AI agents trying to design these "digital soups" (molecules) work like a chef who tastes the soup, guesses what to add, stirs it in, and then tastes it again.

• The Mistake: If they add a heavy spice to boost flavor, the soup might become too salty or hard to digest. If they add a healthy vegetable to make it lighter, it might lose its flavor.
• The Result: The AI keeps making changes, but it rarely improves both flavor and health at the same time. It often fixes one problem while accidentally creating a new one.

The New Idea: "Probe Before You Edit"

The authors, inspired by real-life medicinal chemists, propose a new method called PROBE. Instead of guessing, the AI takes a "test drive" before making any permanent changes.

Think of it like a car mechanic fixing a complex engine.

• Old Way (Current AI): The mechanic says, "I think this part is broken. I'll swap it out!" Then they swap it, start the car, and hope it runs better. If the car stalls, they try something else.
• PROBE Way: The mechanic says, "Before I swap the part, let me run a few controlled tests." They might try adding a tiny bit of fuel, then a little more, then removing a bolt, just to see how the engine reacts to each small change.

How PROBE Works (Step-by-Step)

1. The Map (Site Map Construction)
First, PROBE looks at the molecule and the protein "pocket" it needs to fit into. It breaks the molecule down into small Lego blocks (fragments). It creates a Map that labels each block:

• Synergy Zone: A spot where changing the block will likely make the molecule both stick better and be healthier.
• Tension Zone: A spot where making it stick better will likely make it harder to make (a trade-off).
• Liability Zone: A spot that is dangerous or messy (like a toxic chemical group) that must be fixed, no matter what.

2. The Test Drive (Probing)
Before the AI starts the real optimization, it runs a "simulation test drive." It takes the blocks identified on the Map and makes 12 tiny, controlled changes (probes) to see what happens.

• Example: "If I add a small ring here, does the flavor (binding) go up? Does the health score (druggability) go down?"
• It tests high, medium, and low intensity changes, plus a "reverse" change to see if the original idea was actually correct.

3. The Instruction Manual (EditManual)
After the tests, PROBE writes a Pocket-Specific Instruction Manual. This isn't a generic rulebook; it's a cheat sheet for this specific molecule.

• It says: "For the 'Synergy Zone' block, adding a ring works great. But for the 'Liability Zone' block, never add anything heavy."
• This manual turns the test results into strict rules for the AI to follow.

4. The Team Huddle (Multi-Agent Optimization)
Now, the real optimization begins. PROBE uses a team of three AI agents:

• The Flavor Agent: Wants to make the molecule stick tighter.
• The Health Agent: Wants to make the molecule safer and easier to build.
• The Manager Agent: Listens to both, checks the Instruction Manual, and combines their ideas.
• The Magic: Because they are all following the same manual based on the "test drive" results, they don't fight each other. They find a way to improve the flavor without ruining the health, because the manual told them exactly which moves are safe.

The Results

The researchers tested this on a standard database of drug targets (CrossDocked2020).

• Old AI Agents: Often improved the binding but hurt the safety, or vice versa. They were like a chef who keeps burning the soup while trying to fix the seasoning.
• PROBE: Achieved the best results ever recorded on this test. It successfully improved both binding and safety simultaneously much more often than any other method.

Why This Matters

The paper argues that the biggest mistake current AI makes is committing to a change before knowing how the system will react. By forcing the AI to "probe" (test) first and create a specific guide (the Manual), it stops guessing and starts making informed, safe improvements.

In short: Don't just edit the molecule. Test the edit first, write down the rules, and then edit. That is the secret to better drug design.

Technical Summary

Technical Summary: Probe Before You Edit (PROBE)

1. Problem Statement

Structure-based drug design (SBDD) aims to generate ligands with 3D geometries complementary to specific protein binding pockets. While recent deep generative models (autoregressive, diffusion-based, and language-model-based) have advanced the initial generation of candidate molecules, they typically treat SBDD as a one-time conditional generation task. These models lack mechanisms for iterative inspection, critique, and repair based on pocket-specific feedback.

Conversely, existing approaches employing Large Language Model (LLM) agents for iterative optimization suffer from a critical bottleneck: objective interference. A viable ligand must simultaneously satisfy high binding affinity (measured by Vina scores) and high druggability (measured by QED and Synthetic Accessibility, SA). These objectives often pull in opposite directions under local edits (e.g., adding a hydrophobic group to fill a pocket may improve affinity but degrade druggability by increasing molecular complexity).

Current LLM-agent pipelines (e.g., MoLLM, CIDD, LIDDIA) commit to edits before gathering pocket-specific evidence on how the ligand-pocket complex will respond. Consequently, agents frequently fail to achieve joint improvement (simultaneously improving both affinity and druggability) in a single step, often trading one objective for the other.

2. Methodology: PROBE Framework

The authors propose PROBE (Probing-Guided Molecular Optimization), a framework inspired by the medicinal chemistry practice of "design-make-test-analyze" (DMTA). Unlike prior agents that "edit before knowing," PROBE adopts a "probe-before-edit" strategy. The framework consists of three sequential stages:

A. Site Map Construction

Before optimization begins, PROBE decomposes the initial ligand into fragments (using BRICS) and aggregates three diagnostic signals for each fragment:

1. Interaction Signals: Derived from PLIP (Protein-Ligand Interaction Profiler) regarding hydrogen bonds, hydrophobic contacts, etc.
2. Geometric Signals: Assessing steric clashes, voids, and solvent exposure.
3. Property Signals: Evaluating ligand efficiency, synthetic accessibility, and structural alerts (e.g., PAINS, excessive chirality).

A diagnosis LLM synthesizes these signals to construct a Site Map, labeling specific fragments as:

• SYNERGY: Sites where a single edit can plausibly improve both affinity and druggability.
• TENSION: Sites where improving one objective likely degrades the other.
• LIABILITY: Sites containing structural flaws (e.g., reactive groups, complex spiro systems) that must be fixed regardless of trade-offs.

B. Pocket-Specific Edit Manual Induction

PROBE performs a one-time probing stage to estimate how the specific pocket responds to edits before the iterative loop starts.

1. Probing Planner: An LLM proposes three strategies (Affinity-first, Druggability-first, Co-optimization) targeting the sites identified in the Site Map.
2. Probing Generator: For each strategy, the system generates 12 controlled "probe edits" (3 intensity levels: High, Medium, Low, plus a Counterfactual) per molecule.
3. Response Summarization: The outcomes of these probes (changes in Vina, QED, SA) are analyzed to determine response shapes (e.g., monotone improvement, activity cliffs, saturation).
4. EditManual Construction: The results are distilled into a structured EditManual. This document acts as a pocket-specific guide, defining allowed/forbidden actions, semantic constraints (size, polarity, flexibility), and blacklists for each site. It converts local probe responses into executable constraints for the optimization loop.

C. Multi-Agent Optimization Loop

PROBE runs an iterative DMTA loop guided by the Site Map and EditManual:

1. Design: Three specialized agents (Affinity, Druggability, and Co-optimization) propose edits. The Affinity agent targets SYNERGY/Tension sites for binding; the Druggability agent targets LIABILITY sites for simplification; the Co-optimization agent reconciles proposals.
2. Make & Test: A fragment-assembly engine instantiates the proposed edits using a library of drug-like fragments (ChEMBL-derived) that satisfy the EditManual's semantic constraints. The candidates are scored.
3. Analyze & Iterate: The system updates the EditManual incrementally if outcomes contradict existing rules (e.g., tightening constraints on failed actions). The best molecule (based on hypervolume improvement relative to the initial state) is selected for the next round.

3. Key Contributions

1. Diagnostic Metrics: The authors introduce two metrics to quantify the difficulty of joint optimization:

   • Joint Improvement Rate: The frequency of steps improving both affinity and druggability.
   • Objective Interference (OI) Rate: The frequency where improving one objective causes degradation in the other.
   • Finding: Existing pipelines show low joint improvement rates (<33%) and high OI rates (>60%).

2. PROBE Framework: A novel "probe-before-edit" architecture that explicitly estimates pocket-specific edit responses. It generates a Site Map and an EditManual to guide multi-agent optimization, decoupling exploration from optimization.

3. State-of-the-Art Performance: Evaluation on the CrossDocked2020 benchmark demonstrates that PROBE significantly outperforms existing LLM-agent baselines (MoLLM, CIDD, LIDDIA) and deep generative models across all standard SBDD metrics.

4. Experimental Results

• Benchmark: CrossDocked2020 (100 protein pockets).
• Baselines: Compared against de novo 3D generation models (e.g., Pocket2Mol, MolCRAFT) and LLM-agent optimizers (MoLLM, CIDD, LIDDIA).
• Performance:
  • PROBE achieves the highest Success Ratio (67.67% with MolJO initialization), significantly outperforming the best baseline (MolJO + CIDD at 47.80%).
  • It achieves the best Vina (-9.626), QED (0.664), and SA (0.782) scores among all methods.
  • Joint Improvement: PROBE is the only method where the majority of optimization steps (52.8%) result in joint improvement, compared to <30% for baselines.
  • Objective Interference: PROBE drastically reduces the OI rate (e.g., ~25-30% vs. >60% for baselines), proving that the EditManual successfully constrains edits to avoid collateral damage.
• Ablation Studies:
  • Removing the Site Map (using flat fragments) reduces performance.
  • Removing Probing (relying on LLM prior knowledge for the EditManual) significantly degrades results, confirming that the value comes from evidence-based constraints, not just the existence of a manual.
• Compute Budget: Even when CIDD is extended to match PROBE's computational budget (27 rounds vs. 5), PROBE yields superior results, indicating efficiency gains from the probing strategy rather than just increased compute.

5. Significance and Claims

The paper claims that the primary bottleneck in current LLM-agent SBDD is not the choice of agent paradigm (planner vs. generator) but the lack of pocket-specific evidence before committing to an edit. By mimicking the medicinal chemist's practice of running controlled analog edits to gauge response, PROBE enables agents to navigate the affinity-druggability trade-off more effectively.

The authors position probe-before-edit as a generalizable pipeline for LLM-driven molecular optimization. They suggest that gathering task-specific evidence before committing to edits could extend beyond SBDD to other multi-objective design settings where objectives are often in tension. The work demonstrates that explicit estimation of edit responses, distilled into structured constraints, allows LLM agents to achieve joint optimization outcomes that were previously unattainable with blind iterative editing.