Amortized Molecular Optimization via Group Relative Policy Optimization

TL;DR

AMORTIX is a novel amortized Graph Transformer model for efficient, constraint-aware molecular optimization that outperforms existing methods and generalizes well to unseen structures.

Contribution

The paper introduces AMORTIX, a graph transformer model that supports structural constraints and stabilizes training through reward normalization, enabling efficient molecular optimization.

Findings

AMORTIX outperforms baseline methods on goal-directed scaffold decoration.

It ranks first among amortized methods on the PMO benchmark.

The model successfully transfers learned modifications to unseen drug structures.

Abstract

In structurally constrained molecular optimization, state-of-the-art methods restart an expensive oracle-driven search from scratch for every new input structure, 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, an amortized Graph Transformer model that natively supports such constraints, optimizing molecular structures in a single forward pass 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 show that, under this heterogeneity, standard reinforcement learning methods fail to stabilize training, and address this by normalizing rewards within groups of completions sharing the same starting structure. We evaluate on structurally constrained single- and multi-target kinase inhibitor design, and on a few-shot prodrug case study. AMORTIX outperforms both amortized and instance-optimization baselines on goal-directed scaffold decoration and ranks first among amortized methods on the PMO benchmark; the prodrug case study further demonstrates transfer of a learned modification rule to unseen drug structures. Code is available at https://github.com/Hash-hh/AMORTIX/.