ReMedI: Inference-Time Refinement and Self-Distillation for Controllable Molecular Optimization

Jia Yi Zhang; Shivanshu Shekhar; Allison McLevige; Tong Zhang

cs.LGOct 23, 2026

ReMedI: Inference-Time Refinement and Self-Distillation for Controllable Molecular Optimization

Authors: Jia Yi Zhang, Shivanshu Shekhar, Allison McLevige, Tong Zhang

Organizations: University of Illinois Urbana-Champaign

Abstract

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	$q_{0}$ -base	$q_{0}$ -MH	$q_{1}$ -base	$q_{1}$ -MH
GeLLMO-Llama	SOR $\uparrow$	IND	0.296	0.368	0.630	0.486	0.696
	SOR $\uparrow$	OOD	0.316	0.305	0.506	0.390	0.548
	SSOR $\uparrow$	IND	0.054	0.074	0.170	0.130	0.244
	SSOR $\uparrow$	OOD	0.186	0.163	0.251	0.230	0.286
GeLLMO-Mistral	SOR $\uparrow$	IND	0.530	0.354	0.590	0.472	0.706
	SOR $\uparrow$	OOD	0.390	0.271	0.467	0.332	0.552

Table 1: C-MuMOInstruct results on 100 held-out leads per task.

q_{0}

denotes the original proposal and

q_{1}

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 $\uparrow$		Sim $\uparrow$		CLR $\uparrow$
	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.

Panel of Figure 5 — Figure 5: Effect of repeated proposal distillation. Red dashed lines show SOR/SSOR; blue lines show mean lead similarity. $q_{0}$ is the original proposal and $q_{1}$ – $q_{3}$ are successive distillation rounds.

Figure 6: Effect of distillation sources on $q_{1}$ -Base and $q_{1}$ -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	$q_{0}$ -base	$q_{0}$ -MH	$q_{1}$ -base	$q_{1}$ -MH
GeLLMO-Llama	SOR $\uparrow$	IND	0.320	0.340	0.632	0.532	0.756
	SOR $\uparrow$	OOD	0.290	0.300	0.448	0.424	0.524
	SSOR $\uparrow$	IND	0.030	0.092	0.200	0.140	0.264
	SSOR $\uparrow$	OOD	0.140	0.180	0.276	0.240	0.332
	RI (diag.)	IND	0.61	0.23	0.64	0.49	1.14
	RI (diag.)	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

q_{0}

is the frozen proposal,

q_{1}

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

q_{0}

-base .

Chains $\times$ steps	SOR $\uparrow$	SSOR $\uparrow$	Sim $\uparrow$
$80\times 1$	0.440	0.070	0.463
$40\times 2$	0.500	0.100	0.422
$20\times 4$	0.560	0.125	0.372
$10\times 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

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

See all 103 papers

C-MORAL: Controllable Multi-Objective Molecular Optimization with Reinforcement Alignment for LLMs
Apr 24, 2026Rui Gao, Youngseung Jeon, Swastik Roy +2RL for Language Models Molecular Optimization
University of California, Los Angeles · Amazon
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
MolMem: Memory-Augmented Agentic Reinforcement Learning for Sample-Efficient Molecular Optimization
Apr 14, 2026Ziqing Wang, Yibo Wen, Abhishek Pandy +2Molecular Optimization Sample-Efficient RL
Northwestern University · AbbVie

Apr 24, 2026J/K move · Enter open · S save

C-MORAL: Controllable Multi-Objective Molecular Optimization with Reinforcement Alignment for LLMs

Rui Gao, Youngseung Jeon, Swastik Roy, Morteza Ziyadi+1

University of California, Los Angeles · Amazon

RL for Language Models Molecular Optimization Large Language Model-Guided Optimization Multi-Objective Language Model Alignment Drug Discovery

Open paper AskPDF arXiv

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 S

^2

-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 S

^2

-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.

Abstract

Molecular Optimization

Explore similar work

C-MORAL: Controllable Multi-Objective Molecular Optimization with Reinforcement Alignment for LLMs

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

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