Controllable molecular optimization improves drug lead properties efficiently
ReMedI: Inference-Time Refinement and Self-Distillation for Controllable Molecular Optimization
Machine LearningArtificial Intelligence
Summary
Improving drug molecules is tricky because changes can help some properties but harm others. The authors created a method called ReMedI that refines molecular designs step-by-step during inference time, making better improvements without needing lots of new training data. They use a smart search to tweak molecules while keeping them similar to the original and meeting multiple goals. This method works better than previous systems and learns to generate improved molecules more efficiently afterward.
What this means in practice
- •For pharmaceutical chemists: Optimize drug candidates for multiple desired properties simultaneously while preserving core chemical structures more efficiently than existing methods.
- •For chemical informatics teams: Reduce the amount of specialized data needed to improve molecular design systems through inference-time refinement and self-distillation.
Authors
Jia Yi Zhang, Shivanshu Shekhar, Allison McLevige, Tong Zhang
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.