In drug development, lead optimization must improve competing properties, such as increasing permeability and reducing toxicity, while limiting scaffold changes that could compromise the lead's potency. RL-based post-training, such as C-MORAL, addresses this problem by fine-tuning molecular language models to directly optimize molecular leads for objectives of interest, which relies on expensive task-specific datasets and faces challenges in achieving multi-objective optimization (MOO). We argue a pre-trained model may already contain useful molecular edits that zero-shot generation cannot efficiently reach. Motivated by the limited search capabilities of zero-shot generation, we develop an efficient Metropolis-Hastings search strategy to better explore the proposal distribution. We propose ReMedI, which uses the molecular language model as a base proposer for inference-time alignment. Given a lead and a set of multi-objective tasks, the model proposes edits conditioned on the current molecule, while Metropolis--Hastings refinement guides sequential edits under an objective for property improvement, preservation, and lead similarity. On both C-MuMOInstruct and SMDD benchmarks, refinement improves frozen base models across three backbones, reaching significantly higher multi-objective optimization success rates compared to baselines. These search traces can then distilled into the base model which improves zero-shot generation and subsequent search using approximately 7.8 times less training molecular leads than that of C-MORAL. ReMedI achieves approximately 2.8 times the strict-success coverage of RL baseline C-MORAL, including on out-of-distribution tasks, with matched lead-relative similarity.
Figures & tables
Figure 1: ReMedI separates the molecular edit proposal from the inference-time objective. Refinement discovers successful trajectories; lead- and state-conditioned examples from those trajectories train a shared proposal for later generation and search.
Backbone
Metric
Split
C-MORAL
q0 -base
q0 -MH
q1 -base
q1 -MH
GeLLMO-Llama
SOR ↑
IND
0.296
0.368
0.630
0.486
0.696
OOD
0.316
0.305
0.506
0.390
0.548
SSOR ↑
IND
0.054
0.074
0.170
0.130
0.244
OOD
0.186
0.163
0.251
0.230
0.286
GeLLMO-Mistral
SOR ↑
IND
0.530
0.354
0.590
0.472
0.706
OOD
0.390
0.271
0.467
0.332
0.552
Table 1: C-MuMOInstruct results on 100 held-out leads per task. q0 denotes the original proposal and q1 the proposal after one distillation round; base denotes direct generation and MH refinement. Best-of- 20 selection prioritizes SOR, then SSOR, then similarity. Success–similarity tradeoffs are shown in Fig. 2 . Bold marks the best value in each row.
Figure 2: Strict-success coverage versus minimum lead similarity on 25 matched leads per task. Each method selects candidates by SOR success, then maximum similarity; coverage counts selected outputs that also satisfy SSOR and the indicated similarity threshold. C-MORAL OOD results are averaged over five task-specific GRPO policies.
Backbone
SOR ↑
Sim ↑
CLR ↑
Ours
C-MORAL
Ours
C-MORAL
Ours
C-MORAL
GeLLMO-Llama
0.764
0.580
0.633
0.707
0.764
0.583
GeLLMO-Mistral
0.771
0.640
0.663
0.737
0.792
0.681
mCLM-4B
0.796
0.506
0.316
0.361
0.806
0.576
Table 2: SMDD-Bench Task 4 results on 314 instances. ReMedI-MH uses the frozen proposal, and pooled C-MORAL uses one shared adapter per backbone at the same 20-candidate inference budget.
Figure 3: Evobrutinib case study. Starting from the same “fallen angel” drug lead Edwards et al. (2026) , ReMedI identifies a candidate that simultaneously reduces all three predicted liabilities while preserving substantial structural similarity. This example illustrates how inference-time refinement can incorporate multiple beneficial edits into a single molecule while retaining the lead scaffold.
Figure 4: SOR versus candidate budget on 250 GeLLMO-Mistral instances spanning ten C-MuMOInstruct tasks. At matched candidate budgets, MH achieves the highest observed SOR. Legend values report mean lead similarity.
Figure 5: Effect of repeated proposal distillation. Red dashed lines show SOR/SSOR; blue lines show mean lead similarity. q0 is the original proposal and q1 – q3 are successive distillation rounds.
Figure 6: Effect of distillation sources on q1 -Base and q1 -MH. − A, − B, and − C remove successful endpoints, intermediate states, and improving transitions, respectively. Blue shows mean similarity; red dashed shows SSOR.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Backbone
Metric
Split
C-MORAL
q0 -base
q0 -MH
q1 -base
q1 -MH
GeLLMO-Llama
SOR ↑
IND
0.320
0.340
0.632
0.532
0.756
OOD
0.290
0.300
0.448
0.424
0.524
SSOR ↑
IND
0.030
0.092
0.200
0.140
0.264
OOD
0.140
0.180
0.276
0.240
0.332
RI (diag.)
IND
0.61
0.23
0.64
0.49
1.14
OOD
0.52
0.31
0.51
0.49
0.67
Appendix
Table 3: C-MuMOInstruct performance of C-MORAL and ReMedI across GeLLMO-C (Llama, Mistral) and mCLM-4B backbones. n=50 held-out seeds/task, best-of- 20 . q0 is the frozen proposal, q1 one round of search distillation, base direct generation, and MH inference-time refinement. This 50-lead-per-task summary uses the candidate selection in the source draft (success, then strict success, then similarity), distinct from the success-then-maximum-similarity selector of the 25-lead coverage figures. SOR, SSOR, and Sim are reported across methods; RI is a within-method diagnostic and is not comparable across methods (C-MORAL’s RI is computed on a different scale), so RI rows are not bolded. Bold marks the best comparable value in each row, excluding q0 -base .
Chains × steps
SOR ↑
SSOR ↑
Sim ↑
80×1
0.440
0.070
0.463
40×2
0.500
0.100
0.422
20×4
0.560
0.125
0.372
10×8
0.550
0.140
0.343
Appendix
Table 4: Effect of trajectory depth on frozen GeLLMO-Mistral. Every configuration uses 80 proposals per lead. Results aggregate IND and OOD; bold marks the best value in each column.
Molecular optimization computationally designs or edits chemical structures to improve desired properties while respecting chemical and practical constraints.
Large language models (LLMs) show promise for molecular optimization, but aligning them with selective and competing drug-design constraints remains challenging. We propose C-Moral, a reinforcement learning post-training framework for controllable multi-objective molecular optimization. C-Moral combines group-based relative optimization, property score alignment for heterogeneous objectives, and bottleneck-sensitive non-linear reward aggregation to improve stability across competing molecular properties. Experiments on C-MuMOInstruct and S2-Bench MolOpt show that C-Moral achieves the best performance among compared methods on both benchmarks. On C-MuMOInstruct, C-Moral achieves the best Success Optimized Rate (SOR) of 48.9% on in-domain tasks and 39.5% on out-of-domain tasks while preserving scaffold similarity. On S2-Bench MolOpt, it also achieves the strongest results across LogP, MR, and QED optimization tasks. These results suggest that C-Moral is an effective way to align molecular LLMs with continuous and constrained molecular design objectives. Our code and models are publicly available at https://github.com/Rwigie/C-MORAL.
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 × 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.
Xuan Li, Zhanke Zhou, Zongze Li +4
Hong Kong Baptist University · CMIC, Shanghai Jiao Tong University · DAMO Academy, Alibaba Group +1
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× 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.