Amortized Molecular Optimization via Group Relative Policy Optimization

원문 초록

In structurally constrained molecular optimization, state-of-the-art methods treat every new input structure as a separate instance, restarting an expensive oracle-driven search from scratch and scaling poorly to settings with many starting structures or expensive oracles. While amortized approaches that learn a transferable policy could in principle remove this bottleneck, existing methods struggle to generalize to diverse structural constraints at inference time. We present AMORTIX, a Graph Transformer model that natively supports such constraints, optimizing molecular structures in a single rollout with zero inference-time oracle calls. A central challenge for amortized training in this domain is that optimization difficulty varies drastically across starting structures. We address this by centering rewards within groups of completions that share the same starting structure. We evaluate on goal-directed scaffold decoration, including single- and multi-target kinase inhibitor design, and on a few-shot prodrug case study. AMORTIX outperforms both amortized baselines and search-based instance optimizers on kinase scaffold decoration, and ablations across 23 objectives show that group-relative centering improves over REINFORCE and PPO and is substantially more consistent across scaffold splits. The prodrug case study further demonstrates transfer of learned modification rules to unseen drug structures.