MolMem: Memory-Augmented Agentic Reinforcement Learning for Sample-Efficient Molecular Optimization

Ziqing Wang; Yibo Wen; Abhishek Pandy; Han Liu; Kaize Ding

cs.LGApr 14, 2026

MolMem: Memory-Augmented Agentic Reinforcement Learning for Sample-Efficient Molecular Optimization

Authors: Ziqing Wang, Yibo Wen, Abhishek Pandy, Han Liu, Kaize Ding

Organizations: Northwestern University · AbbVie

Abstract

In drug discovery, molecular optimization aims to iteratively refine a lead compound to improve molecular properties while preserving structural similarity to the original molecule. However, each oracle evaluation is expensive, making sample efficiency a key challenge for existing methods under a limited oracle budget. Trial-and-error approaches require many oracle calls, while methods that leverage external knowledge tend to reuse familiar templates and struggle on challenging objectives. A key missing piece is long-term memory that can ground decisions and provide reusable insights for future optimizations. To address this, we present MolMem (Molecular optimization with Memory), a multi-turn agentic reinforcement learning (RL) framework with a dual-memory system. Specifically, MolMem uses Static Exemplar Memory to retrieve relevant exemplars for cold-start grounding, and Evolving Skill Memory to distill successful trajectories into reusable strategies. Built on this memory-augmented formulation, we train the policy with dense step-wise rewards, turning costly rollouts into long-term knowledge that improves future optimization. Extensive experiments show that MolMem achieves 90% success on single-property tasks (1.5 $\times$ over the best baseline) and 52% on multi-property tasks using only 500 oracle calls. Our code is available at https://github.com/REAL-Lab-NU/MolMem.

Figures & tables

Figure 1: The Evolution of Molecular Optimization Paradigms. (1) Novice : trial-and-error exploration with low sample efficiency; (2) Apprentice : template imitation of external knowledge; (3) Expert (Ours) : self-evolution via grounded knowledge and consolidated skills.

Figure 2: Overview of the MolMem . The framework consists of two complementary memory modules: (Left) Static Exemplar Memory retrieves structurally similar, high-scoring molecules from an external database as references for the current objective. (Right) Evolving Skill Memory distills successful trajectories into reusable textual strategies. Retrieved exemplars or skills are injected into the Working Memory (center) to augment the LLM agent’s context during Multi-turn Rollouts .

Figure 3: Mechanism of the Evolving Skill Memory. (Top) Memory Formation : High-reward pairs in multi-turn rollouts are distilled into textual skills by extracting structural changes and synthesizing insights via a summarizer agent. (Bottom) Memory Retrieval : Relevant skills are retrieved from the skill bank using a hybrid matching of fingerprint similarity and functional group overlap to guide the policy agent.

❶ Novice
Method	Backbone	QED			plogP			SA			DRD2			JNK3
Method	Backbone	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI
Graph-GA	N/A	59.5	0.49	0.13	61.5	0.49	9.64	46.0	0.48	0.20	34.0	0.56	5.13	6.5	0.59	1.73
QMO	GRU	16.5	0.53	0.16	8.0	0.55	5.87	6.5	0.52	0.15	9.5	0.51	3.91	0.5	0.56	1.20
Reinvent 4	Transformer	33.5	0.50	0.14	51.5	0.48	10.37	18.5	0.47	0.21	34.5	0.50	8.71	28.0	0.51	1.91
❷ Apprentice

Table 1: Single-Property Optimization Results. Methods are categorized by optimization paradigm: Novice , Apprentice , and Expert . SR (%) denotes the success rate, Sim denotes the Tanimoto similarity to the original molecule (we require

\mathrm{Sim}\geq 0.4

), and RI denotes relative improvement. For SR, gold = best, silver = second, and bronze = third.

❶ Novice
Method	Backbone	QED+plogP			plogP+DRD2			QED+SA			DRD2+SA			DRD2+QED+plogP
Method	Backbone	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI	SR (%)	Sim	RI
Graph-GA	N/A	8.0	0.50	4.11	2.5	0.48	6.61	8.0	0.52	0.12	0.0	0.53	4.15	0.0	0.51	3.79
QMO	GRU	3.0	0.52	6.40	2.0	0.51	4.20	1.5	0.53	0.11	1.0	0.50	2.49	0.0	0.52	1.14
Reinvent 4	Transformer	6.0	0.49	8.60	50.0	0.48	7.94	17.0	0.66	0.05	37.5	0.50	5.54	3.0	0.51	4.28
❷ Apprentice

Table 2: Multi-Property Optimization Results. Methods are categorized by optimization paradigm: Novice , Apprentice , and Expert . SR (%) denotes the success rate, Sim denotes the Tanimoto similarity to the original molecule (we require

\mathrm{Sim}\geq 0.4

), and RI denotes relative improvement. For SR (Success Rate), gold = best, silver = second, bronze = third.

Panel of Figure 4 — Figure 4: Analysis of key hyperparameters. (a) Success rate improves with more turns and saturates around $T{=}5$ . (b) Performance improves with more training leads, though even 64 leads yield competitive results. (c) Stricter similarity constraints reduce the success rate.

Panel of Figure 5 — Figure 5: What the skill bank learns (QED task). (a) Distribution of scaffold modifications, where ring removal is the most frequent operation. (b) Functional-group flux suggests a tendency to remove amine groups while adding halogen substitutions. (c) Optimization trajectories indicate that QED improvements often coincide with reductions in molecular weight.

Model Configuration	QED			DRD2
Model Configuration	SR (%)	Sim	RI	SR (%)	Sim	RI
MolMem (Full)	91.0	0.47	0.24	96.0	0.49	16.97
— Ablation of Training Components —
w/o SFT Initialization	68.0	0.54	0.19	34.5	0.58	8.64
w/o Multi-turn (T=1)	50.0	0.55	0.19	54.0	0.58	11.39
w/o RL Training (SFT-only)	37.0	0.62	0.14	20.0	0.62	7.04

Table 3: Ablation Study of Component Contributions.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Property	Single-property success	Multi-property success
QED	$F_{\text{QED}}(m^{\prime})\geq 0.9$	$\Delta F_{\text{QED}}\geq 0.1$
plogP	$F_{\text{plogP}}(m^{\prime})\geq 2.0$	$\Delta F_{\text{plogP}}\geq 1.0$
JNK3	$F_{\text{JNK3}}(m^{\prime})\geq 0.1$	$\Delta F_{\text{JNK3}}\geq 0.1$
DRD2	$F_{\text{DRD2}}(m^{\prime})\geq 0.8$	$\Delta F_{\text{DRD2}}\geq 0.5$
SA	$F_{\text{SA}}(m^{\prime})\leq-2.5$	$\Delta F_{\text{SA}}\leq-0.5$

Appendix

Table 4: Task-specific success criteria. For multi-property tasks,

\Delta F=F(m^{\prime})-F(m)

is measured relative to the lead molecule

m

Category	Parameter	Value
Backbone & SFT	Base model	Qwen2.5-1.5B-Instruct
	SFT epochs	10
	LoRA rank $r$	16
	LoRA $\alpha$	32
PPO	Training steps	100
	PPO learning rate	$5\times 10^{-5}$

Appendix

Table 5: Training hyperparameters for MolMem .

Category	Parameter	Value
Budget & constraints	Oracle-call budget per lead $B$	500
	Similarity threshold $\gamma$	0.4
	Search iterations (generations) $G$	20
Candidate generation	Rollouts per iteration $N$	32
Candidate generation	Max turns per rollout $T$	5
Decoding diversity	Base temperature $\tau_{0}$	0.9

Appendix

Table 6: Inference hyperparameters for MolMem .

Method	Model	One-time	Per-task	Per-task
	Size	Training	Training	Inference
Reinvent 4*	–	–	14h (online)	–
GeLLM 3 O	8B	48h (SFT)	–	1h
MolMem	1.5B	10h (SFT)	4h	3h

Appendix

Table 7: Computational cost comparison. GPU hours are reported on 2

\times

H100. *Online RL baseline.

Panel of Figure 6 — Figure 6: Case studies of learned skills. (Top) DRD2 examples highlighting tail substitution and motif augmentation. (Bottom) QED examples highlighting bioisosteric replacement and side-chain simplification. Each card shows the distilled strategy, the before/after edit, and the corresponding score change.

Figure 7: Memory-triggered case studies. (Top) A QED run where an early proposal violates the similarity constraint; after the agent plateaus, retrieved guidance leads to a feasible edit that improves QED. (Bottom) A QED run where exemplar hints provide a viable direction for subsequent edits. Each panel reports the lead molecule, representative intermediate proposals, and the resulting QED and similarity to the lead.

Molecular Optimization

Molecular optimization computationally designs or edits chemical structures to improve desired properties while respecting chemical and practical constraints.

See all 103 papers

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning
Mar 6, 2026Xuan Li, Zhanke Zhou, Zongze Li +4Molecular Optimization Large Language Model-Guided Optimization
Hong Kong Baptist University · CMIC, Shanghai Jiao Tong University · DAMO Academy, Alibaba Group +1
Small Molecule Optimization with Large Language Models
Jul 26, 2024Philipp Guevorguian, Menua Bedrosian, Tigran Fahradyan +3Molecular Optimization Evolutionary Optimization
YerevaNN research lab, Yerevan, Armenia.
MFG: Multi-turn Feedback Guided Molecular Optimization
Oct 13, 2026Trieu Nguyen, Frazier N. Baker, Ruoxi Gao +2Molecular Optimization Large Language Model-Guided Optimization
Department of Computer Science and Engineering, The Ohio State University · Division of Medicinal Chemistry and Pharmacognosy, The Ohio State University · Department of Biomedical Informatics, The Ohio State University +1

Mar 6, 2026J/K move · Enter open · S save

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Xuan Li, Zhanke Zhou, Zongze Li, Jiangchao Yao+3

Hong Kong Baptist University · CMIC, Shanghai Jiao Tong University · DAMO Academy, Alibaba Group+1

Molecular Optimization Large Language Model-Guided Optimization Policy Optimization RL for Language Models Reinforcement Learning with Verifiable Rewards

Open paper AskPDF arXiv

Large language models (LLMs) benefit substantially from supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR) in reasoning tasks. However, these recipes perform poorly in instruction-based molecular optimization, where each data point typically provides only a single optimized reference molecule and no step-by-step optimization trajectory. We reveal that answer-only SFT on the reference molecules collapses reasoning, and RLVR provides sparse feedback under similarity constraints due to the model's lack of effective exploration, which slows learning and limits optimization. To encourage the exploration of new molecules while balancing the exploitation of the reference molecules, we introduce Reference-guided Policy Optimization (RePO), an optimization approach that learns from reference molecules without requiring trajectory data. At each update, RePO samples candidate molecules with their intermediate reasoning trajectories from the model and trains the model using verifiable rewards that measure property satisfaction under similarity constraints in an RL manner. Meanwhile, it applies reference guidance by keeping the policy's intermediate reasoning trajectory as context and training only the answer in a supervised manner. Together, the RL term promotes exploration, while the guidance term mitigates reward sparsity and stabilizes training by grounding outputs to references when many valid molecular edits exist. Across molecular optimization benchmarks, RePO consistently outperforms SFT and RLVR baselines (e.g., GRPO), achieving improvements on the optimization metric (Success Rate

\times

Similarity), improving balance across competing objectives, and generalizing better to unseen instruction styles. Our code is publicly available at https://github.com/tmlr-group/RePO.

Abstract

Molecular Optimization

Explore similar work

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Small Molecule Optimization with Large Language Models

MFG: Multi-turn Feedback Guided Molecular Optimization