# Amortized Molecular Optimization via Group Relative Policy Optimization

arXiv: https://arxiv.org/abs/2602.12162 · HTML: https://mlsanity.com/paper/2602.12162/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 41 of its 63 references that are in the map.
- First posted on arXiv 12 Feb 2026.

## Builds on (most cited first, 30 of 41)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training](https://mlsanity.com/paper/2501.17161/) (2025)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Self-critical Sequence Training for Image Captioning](https://mlsanity.com/paper/1612.00563/) (2016)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)
- [GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation](https://mlsanity.com/paper/2001.09382/) (2020)
- [Molecular De Novo Design through Deep Reinforcement Learning](https://mlsanity.com/paper/1704.07555/) (2017)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development](https://mlsanity.com/paper/2102.09548/) (2021)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [GraphDF: A Discrete Flow Model for Molecular Graph Generation](https://mlsanity.com/paper/2102.01189/) (2021)
- [Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization](https://mlsanity.com/paper/2206.12411/) (2022)
- [MARS: Markov Molecular Sampling for Multi-objective Drug Discovery](https://mlsanity.com/paper/2103.10432/) (2021)
- [Dirichlet Diffusion Score Model for Biological Sequence Generation](https://mlsanity.com/paper/2305.10699/) (2023)
- [Learning to Extend Molecular Scaffolds with Structural Motifs](https://mlsanity.com/paper/2103.03864/) (2021)
- [Hit and Lead Discovery with Explorative RL and Fragment-based Molecule Generation](https://mlsanity.com/paper/2110.01219/) (2021)
- [Molecule Generation by Principal Subgraph Mining and Assembling](https://mlsanity.com/paper/2106.15098/) (2021)
- [GenMol: A Drug Discovery Generalist with Discrete Diffusion](https://mlsanity.com/paper/2501.06158/) (2025)
- [MIMOSA: Multi-constraint Molecule Sampling for Molecule Optimization](https://mlsanity.com/paper/2010.02318/) (2020)
- [Gotta be SAFE: A New Framework for Molecular Design](https://mlsanity.com/paper/2310.10773/) (2023)
- [Retrieval-based Controllable Molecule Generation](https://mlsanity.com/paper/2208.11126/) (2022)
- [Genetic-guided GFlowNets for Sample Efficient Molecular Optimization](https://mlsanity.com/paper/2402.05961/) (2024)
- [A Deep Generative Model for Molecule Optimization via One Fragment Modification](https://mlsanity.com/paper/2012.04231/) (2020)

## Abstract (arXiv)

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