Amortized Molecular Optimization via Group Relative Policy Optimization

arXiv (Cornell University)

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

Authors 8

1. Muhammad bin Javaid
2. Hasham Hussain Aachen Department of Computer Science

   Alfred E. Tiefenbacher (Germany) · RWTH Aachen University

   Affiliation as printed

   Alfred E. Tiefenbacher (GmbH & Co. KG) , Van-der-Smissen-Strasse 1 , 22767 Ham- burg , Germany

   RWTH Aachen University , Department of Computer Science , Ahornstrasse 55 , 52074 Aachen , Germany

3. Ashima Khanna

   Weihenstephan-Triesdorf University of Applied Sciences · Stadtwerke Straubing (Germany)

   Affiliation as printed

   Technical University of Munich , TUM Campus Straubing for Biotechnology and Sustainability, Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

   University of Applied Sciences Weihenstephan-Triesdorf , Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

4. Berke Kisin
5. Jonathan Pirnay

   Weihenstephan-Triesdorf University of Applied Sciences · Stadtwerke Straubing (Germany)

   Affiliation as printed

   Technical University of Munich , TUM Campus Straubing for Biotechnology and Sustainability, Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

   University of Applied Sciences Weihenstephan-Triesdorf , Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

6. Alexander Mitsos Aachen

   RWTH Aachen University

   Affiliation as printed

   RWTH Aachen University , Pro- cess Systems Engineering (AVT.SVT) , Forckenbeckstrasse 51 , 52074 Aachen , Germany

7. Dominik G. Grimm

   Weihenstephan-Triesdorf University of Applied Sciences · Stadtwerke Straubing (Germany)

   Affiliation as printed

   Technical University of Munich , TUM Campus Straubing for Biotechnology and Sustainability, Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

   University of Applied Sciences Weihenstephan-Triesdorf , Bioinformatics , Petersgasse 18 , 94315 Straubing , Germany

8. Martin Grohe Aachen Department of Computer Science

   RWTH Aachen University

   Affiliation as printed

   RWTH Aachen University , Department of Computer Science , Ahornstrasse 55 , 52074 Aachen , Germany