Constraint-Aware Corrective Memory for Language-Based Drug Discovery Agents

Maochen Sun; Youzhi Zhang; Gaofeng Meng

cs.AIApr 10, 2026

Constraint-Aware Corrective Memory for Language-Based Drug Discovery Agents

Authors: Maochen Sun, Youzhi Zhang, Gaofeng Meng

Organizations: Institute of Automation, Chinese Academy of Sciences, China · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences, China · Centre for Artificial Intelligence and Robotics, Hong Kong Institute of Science & Innovation, Chinese Academy of Sciences, China

Abstract

Large language models are making autonomous drug discovery agents increasingly feasible, but reliable success in this setting is not determined by any single action or molecule. It is determined by whether the final returned set jointly satisfies protocol-level requirements such as set size, diversity, binding quality, and developability. This creates a fundamental control problem: the agent plans step by step, while task validity is decided at the level of the whole candidate set. Existing language-based drug discovery systems therefore tend to rely on long raw history and under-specified self-reflection, making failure localization imprecise and planner-facing agent states increasingly noisy. We present CACM (Constraint-Aware Corrective Memory), a language-based drug discovery framework built around precise set-level diagnosis and a concise memory write-back mechanism. CACM introduces protocol auditing and a grounded diagnostician, which jointly analyze multimodal evidence spanning task requirements, pocket context, and candidate-set evidence to localize protocol violations, generate actionable remediation hints, and bias the next action toward the most relevant correction. To keep planning context compact, CACM organizes memory into static, dynamic, and corrective channels and compresses them before write-back, thereby preserving persistent task information while exposing only the most decision-relevant failures. Our experimental results show that CACM improves the target-level success rate by 36.4% over the state-of-the-art baseline. The results show that reliable language-based drug discovery benefits not only from more powerful molecular tools, but also from more precise diagnosis and more economical agent states.

Figures & tables

Figure 1. Overview of constraint-aware corrective memory (CACM) . CACM maintains an agent state composed of static memory, dynamic descriptive memory, and corrective memory. The planner selects one action from Generate , Optimize , and Code/Screen . The executed action produces a candidate molecule set, which is checked by a deterministic protocol audit. If the current set satisfies the returned-set protocol, it is returned as the final output. Otherwise, the audited result is reorganized into corrective memory, compressed, and adapted into the next agent state. CACM therefore turns raw history into a compact protocol-aware agent state.

Metric	Pocket2Mol		DiffSMOL		Claude		GPT-4o		o1-mini		o1		LIDDiA (DeepSeek)		CACM (ours)
Metric	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t	%m/t or %t	#m/t or #t
Generated molecules
Generated	–	100.0	–	100.0	–	5.0	–	5.0	–	5.0	–	5.0	–	21.0	–	5.0
Valid	100.0	100.0	99.9	99.9	98.7	4.9	97.3	4.9	91.3	4.6	95.3	4.8	100.0	21.0	100.0	5.0
QED $\geq\overline{\mathrm{QED}}_{t}$	53.4	53.4	60.0	60.0	96.7	4.8	88.2	4.4	90.1	4.5	88.3	4.4	96.6	20.2	100.0	5.0
LRF $\geq\overline{\mathrm{LRF}}_{t}$	99.7	99.7	72.1	72.1	98.7	4.9	95.9	4.8	90.7	4.5	95.3	4.8	97.0	20.2	100.0	5.0

Table 1. Main results on the 30-target LIDDiA benchmark. Numbers for Pocket2Mol, DiffSMOL, Claude, GPT-4o, o1-mini, and o1 are taken from the original LIDDiA paper; the last two columns are recomputed from our logs.

Variant	Repair Signal	Structured Memory	Dynamic Compression	Corrective Selection	TSR (% / #)	DVS & HQ (% / #)	Avg. Pool Size	Avg. Term. Iters
LIDDiA baseline	✗	✗	✗	✗	73.3 / 22	80.0 / 24	21.0	4.40
Set-level Repair Signal	✓	✗	✗	✗	86.7 / 26	86.7 / 26	6.7	4.07
CACM w/o Corrective Selection	✓	✓	✓	✗	83.3 / 25	83.3 / 25	8.0	4.27
CACM w/o Dynamic Compression	✓	✓	✗	✓	96.7 / 29	96.7 / 29	8.6	3.40
CACM	✓	✓	✓	✓	100.0 / 30	100.0 / 30	5.0	3.07

Table 2. Ablation study of CACM. TSR is computed under the same LIDDiA success definition for all variants. DVS & HQ is computed from the returned pool. Avg. Term. Iters is averaged over 30 targets.

Figure 2. Planner-facing memory length across iterations. Average planner-facing memory length over iterations on the 30-target benchmark, comparing the LIDDiA-style raw-history write-back with the CACM compressed agent state.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Metric	Aggregation	Ref.
Pool size	$\|S\|$	benchmark
Diversity	set-level diversity functional	benchmark
Novelty	set-level novelty functional	benchmark
QED	$\min_{x\in S}\mathrm{QED}(x)$	( Bickerton et al., 2012 )
SAS	$\max_{x\in S}\mathrm{SAS}(x)$	( Ertl and Schuffenhauer, 2009 )
Lipinski check	worst-case pool aggregation	( Lipinski et al., 2001 )

Appendix

Table 3. Summary of the main metrics used in the deterministic audit.

Symbol	Meaning
$K_{d}$	Maximum number of selected pool summaries retained in dynamic descriptive memory.
$W_{d}$	Maximum number of retained recent actions in dynamic descriptive memory.
$K_{c}$	Maximum number of selected corrective entries retained in corrective memory.
$B_{s}$	Character budget for static memory after deterministic template formatting.
$B_{d}$	Character budget for dynamic descriptive memory after deterministic template formatting.
$B_{c}$	Character budget for corrective memory after deterministic template formatting.

Appendix

Table 4. Implementation-side budget variables used to instantiate bounded selection and compression in CACM.

Variable	Value
$K_{d}$	4
$W_{d}$	3
$K_{c}$	3
$B_{s}$	1400
$B_{d}$	1800
$B_{c}$	1200

Appendix

Table 5. Default implementation-side budgets used in CACM. These values are fixed across the full 30-target benchmark without target-specific tuning.

Setting	$(K_{d},W_{d},K_{c})$	$(B_{s},B_{d},B_{c})$
default	$(4,3,3)$	$(1400,1800,1200)$
tight_chars	$(4,3,3)$	$(900,1200,700)$
compact_counts	$(2,2,2)$	$(1400,1800,1200)$
wide_counts	$(6,5,5)$	$(1400,1800,1200)$
rebalanced_chars	$(4,3,3)$	$(1000,2200,1000)$

Appendix

Table 6. Budget configurations used in the CACM sensitivity study.

Setting	TSR (% / #)	Avg. Pool Size	Avg. Term. Iters	Avg. State Chars
default	100.0 / 30	5.00	3.07	1572.2
tight_chars	96.7 / 29	8.17	3.47	1318.0
compact_counts	96.7 / 29	8.90	3.50	1613.4
wide_counts	100.0 / 30	5.53	3.17	1580.0
rebalanced_chars	96.7 / 29	8.17	3.43	1498.4

Appendix

Table 7. Sensitivity of CACM to implementation-side budgets on the 30-target benchmark.

Setting	Static	Dynamic	Corrective
default	601.2	241.2	729.8
tight_chars	601.2	264.8	451.9
compact_counts	601.2	233.0	779.3
wide_counts	601.2	256.9	721.9
rebalanced_chars	601.2	262.8	634.4

Appendix

Table 8. Channel-wise average lengths of the planner-facing agent state in the CACM sensitivity study. Values are averaged over planner-facing states within each trajectory.

Target	Returned molecules
ADRB1	1. Cn1ccc(-c2cc(C(=O)N3CCN(Cc4ccccn4)CC3)c3ccccc3n2)n1 2. O=C(/C=C/c1ccccc1N+[O-])Nc1cccc(Cl)c1 3. Cc1nc(C)c(C@HNC(=O)CCC(=O)c2ccc(Cl)cc2)s1 4. COC(=O)N1CCC@@HC1 5. CCn1c(CCC@@HC=O)nc2ccccc21
ADRB2	1. O=C(CCn1cnc2ccccc2c1=O)OCc1ccc(-c2ccccc2)cc1 2. CC1(c2ccc(OS(=O)(=O)c3ccc(Cl)cc3)cc2)SCCS1 3. Cc1occc1C@HCc1c(F)cccc1F 4. Cc1ccc(C(=O)N[C@H]2CCC(=O)NC2=O)c(N)c1 5. Cc1cc(C(=O)N2CCN(Cc3ccsc3)CC2)no1
AR/NR3C4	1. Cc1ccc(CC@HNC(=O)Nc2ccccc2CC(=O)N(C)C)c(C)c1 2. Fc1cccc(F)c1N[C@@H]1CCOC2(CCCC2)C1 3. Cn1c(Cl)cnc1CS@CCCc1ccc(Cl)cc1 4. CCc1nc(N2CCOCC2)c2onc(-c3ccc(C)cc3)c2n1 5. O=C(NCCOc1ccc(Br)cc1)[C@H]1CCS(=O)(=O)C1
BCHE	1. O=C(NCC(=O)N1CCOc2ccccc21)NC1CCCCC1 2. Cc1ccc(S(=O)(=O)CC@HCOc2ccc(F)cc2Cl)cc1 3. Cc1sc(NC(=O)c2ccon2)nc1-c1ccc2c(c1)CCC2 4. CNC(=O)c1c(C)cccc1NC(=O)Cc1cccc(OCC#N)c1 5. O=S(=O)(NCCCn1cncn1)c1cc(Cl)cc(Cl)c1
CDK5	1. COc1ccc(N2CC@@HCC2=O)cc1 2. N#Cc1ccc(-c2cc(F)cc(C(=O)[O-])c2)cc1O 3. COc1cc(C)c(-c2nc(C3CCCCC3)no2)cc1OC 4. CSc1ccc(Cl)c(C(=O)Nc2cc(C(=O)N(C)C)ccc2Cl)c1 5. CC@@Hc1cccc(-n2cccn2)c1
CHK2	1. Cc1ccc(NS(=O)(=O)c2cc(C(=O)Nc3ccccc3C(=O)[O-])ccc2Br)cc1 2. c1ccc(SCc2noc(-c3ccco3)n2)nc1 3. O=C(c1ccc(OCC2CC2)nc1)N1CCc2[nH]c(=O)sc2C1 4. COc1cc(Br)c(/C=C2\SC(=N)NC2=O)cc1OC 5. COc1ccc(CCNC(=O)[C@H]2CC(=O)N(c3cccc(C(F)(F)F)c3)C2)cc1

Appendix

Table 9. Terminal molecules returned by CACM for all 30 targets. Each row corresponds to one target and lists the five molecules in the final returned set.

Method

Iter

Generated

Valid

QED

LRF

SAS

VNA

NVT

DVS

>0.8

N{\geq}5

&DVS

N{\geq}5

&HQ

DVS&HQ

TSR

CACM (full)

26.1

100.0 / 26.1

95.9 / 22.0

98.6 / 24.7

86.3 / 12.9

79.6 / 13.7

93.9 / 20.0

70.6 / 5.0

30/30

20/30

29/30

18/30

8.2

100.0 / 8.2

99.8 / 7.9

96.3 / 7.0

92.9 / 5.5

99.1 / 7.3

91.5 / 4.7

30/30

27/30

29/30

27/30

11.3

100.0 / 11.3

98.2 / 9.5

99.8 / 11.2

98.2 / 9.6

95.6 / 7.0

96.8 / 8.2

93.7 / 5.0

30/30

29/30

28/30

5.0

100.0 / 5.0

30/30

5.0

100.0 / 5.0

30/30

CACM w/o dynamic compression

18.4

100.0 / 18.4

94.8 / 13.3

99.0 / 17.5

95.0 / 13.4

88.3 / 10.6

95.2 / 13.7

83.3 / 5.6

30/30

27/30

29/30

24/30

Appendix

Table 10. Cutoff statistics at 2/4/6/8/10 iterations for the retained variants in the final ablation.

Method	Planner / target	Extra / target	Extra share	Total / target	Planner / iter	Extra / iter	Total / iter
LIDDiA baseline	24,799.0	8,719.3	26.0%	33,518.3	5,510.9	1,937.6	7,448.5
CACM	9,927.9	13,114.4	56.9%	23,042.2	3,423.4	4,522.2	7,945.6

Appendix

Table 11. Success-truncated controller-side token usage on the 30-target benchmark. For each target, tokens are accumulated only up to the first successful stopping point; if no success is reached, the full 10-iteration budget is used. “Extra” denotes the metric-judgment stage for the LIDDiA baseline and the diagnosis stage for CACM.

Method	Avg. cutoff iterations	Minutes / iter	Minutes / target
LIDDiA baseline	4.40	10.11	45.49
CACM	3.07	15.14	43.91

Appendix

Table 12. Success-truncated wall-clock runtime on the 30-target benchmark. For each target, runtime is accumulated only up to the first successful stopping point; if no success is reached, the full 10-iteration budget is used.

Drug Discovery

Drug discovery identifies, designs, and tests molecules that can safely affect disease biology.

See all 164 papers

DrugSAGE:Self-evolving Agent Experience for Efficient State-of-the-Art Drug Discovery
May 14, 2026Yikun Zhang, Xiwei Cheng, Tianyu Liu +2AI Agents for Scientific Discovery AI-Assisted Scientific Research
1Northeastern University · 2Broad Institute of MIT and Harvard · 3Yale University +1
MADDIE: Multi-Agent Drug Discovery Engine
Oct 12, 2026Ye Liu, Reza Averly, Avery Meyer +4Drug Discovery Multi-Agent LLM Systems
Biomedical Informatics, The Ohio State University. · Computer Science and Engineering, The Ohio State University. · Medicinal Chemistry and Pharmacognosy, The Ohio State University. +1
Mozi: Governed Autonomy for Drug Discovery LLM Agents
Mar 4, 2026He Cao, Siyu Liu, Fan Zhang +7AI Agents for Scientific Discovery LLM Agent Orchestration

May 14, 2026J/K move · Enter open · S save

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

Yikun Zhang, Xiwei Cheng, Tianyu Liu, Yuanqi Du+1

1Northeastern University · 2Broad Institute of MIT and Harvard · 3Yale University+1

AI Agents for Scientific Discovery AI-Assisted Scientific Research Molecular Property Prediction Memory-Augmented Agents Self-Evolving Agents Drug Discovery

Open paper AskPDF arXiv

Building state-of-the-art (SOTA) predictive models for drug discovery requires expensive search over tools, architectures, and training strategies. Current LLM-based agents can find SOTA solutions through extensive trial and error, but they do not retain the experience accumulated along the way and therefore pay the full search cost on every new task. We propose \method (Self-evolving Agent Experience), a framework that accumulates and reuses experience across tasks to build SOTA drug discovery models efficiently. \method maintains a cross-task memory of verified skills, statistical evidence about effective strategies, and a record of recurring errors and their fixes. In some cases, \method transfers a working solution directly without test-time search. In 33 molecular property prediction tasks, \method ranks first among nine SOTA agents in a single-task setting. With memory accumulated from 16 smaller tasks, \method achieves an averaged normalized score of 0.935 on 17 held-out tasks in a cross-task evaluation setting and outperforms all baseline agents by 10-30% in a zero-test-time search regime. In summary, our work shows the advantage of cross-task memory for efficient SOTA model development in drug discovery.

Abstract

Drug Discovery

Explore similar work

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

MADDIE: Multi-Agent Drug Discovery Engine

Mozi: Governed Autonomy for Drug Discovery LLM Agents