Amortized Molecular Optimization via Group Relative Policy Optimization

AI summary

Overview

Research area: Machine learning for molecular design — specifically reinforcement learning (RL) policies for structure-constrained molecule generation in drug discovery.

Technical level: Intermediate. Readers need some familiarity with reinforcement learning (policy gradients, advantage estimation, PPO/REINFORCE) and with graph-based molecular generation, but the core insight is explained in plain terms.

Scope: The paper introduces AMORTIX, a Graph Transformer policy trained with a Group Relative Policy Optimization (GRPO) variant that learns to optimize molecules conditioned on a starting structure in a single forward pass, with no oracle calls at inference time.

What This Paper Is About

Most state-of-the-art molecular optimization methods restart an expensive oracle-driven search from scratch for every new starting structure, which scales poorly when there are many structures or when the scoring oracle is costly. Amortized methods — which learn a transferable policy once and reuse it — could remove that bottleneck, but existing amortized approaches struggle to generalize to diverse structural constraints at inference time. The authors argue that a major cause is the wide variation in optimization difficulty across starting structures, and they propose a graph-based policy trained with group-relative reward normalization to fix it.

Key Contributions

1. Diagnosis of a failure mode. The authors identify high variance in reward signals caused by heterogeneous molecular scaffold difficulty as a primary failure mode for amortized RL in structurally constrained molecular optimization, and show through ablation that neither REINFORCE with a global baseline nor PPO with a learned value critic resolves it.

2. The AMORTIX model. A decoder-only Graph Transformer policy that natively supports structural constraints and generates optimized molecules for unseen inputs in a single forward pass with zero inference-time oracle calls.

3. A GRPO adaptation for molecules. Group Relative Policy Optimization is adapted to the molecular domain: a group of G completions is sampled per starting structure and rewards are normalized against the group mean, decoupling the learning signal from cross-structure difficulty variation, with no critic and no PPO-style clipping.

4. Empirical validation across three settings. AMORTIX outperforms both amortized and instance-optimization baselines on goal-directed scaffold decoration for GSK3β and JNK3 and a Kinase MPO task; it ranks first among amortized methods and second overall on the PMO benchmark; and a prodrug case study demonstrates transfer of a learned modification rule to unseen drug structures.

Main Findings

• Scaffold decoration dominance. On goal-directed scaffold decoration over three test folds, AMORTIX scores 0.775 ± 0.005 on GSK3β, 0.796 ± 0.004 on JNK3, and 0.619 ± 0.004 on Kinase MPO — the highest score on all three tasks, beating both amortized and instance-optimization baselines.

• Amortized baselines lag badly. LibINVENT scores 0.455 ± 0.006 on Kinase MPO versus AMORTIX's 0.619 ± 0.004, while GraphXForm (0.064 ± 0.008 GSK3β, 0.024 ± 0.003 JNK3) and DrugEx v3 (0.194 ± 0.004, 0.356 ± 0.003) score low on the biological targets, suggesting their policies do not move far from the pre-training distribution.

• Instance optimizers win on heuristics, lose on biology. Mol GA (0.461 ± 0.013) and GenMol (0.479 ± 0.012) trail AMORTIX on Kinase MPO. The sub-task breakdown shows Mol GA reaching QED 0.81 and GenMol reaching QED 0.88 with SA 2.93, both scoring below 0.24 on GSK3β/JNK3, whereas AMORTIX prioritizes the difficult biological targets (0.74/0.59) at moderate QED/SA cost.

• Ablations confirm the heterogeneity hypothesis. AMORTIX-DeNovo collapses (0.572 ± 0.015 GSK3β, 0.381 ± 0.029 JNK3, 0.397 ± 0.005 MPO), confirming structural conditioning is needed. AMORTIX-REINFORCE matches AMORTIX on single-objective tasks (0.757 ± 0.085, 0.754 ± 0.092) but is severely unstable on MPO (0.475 ± 0.120, with sub-task standard deviations up to ± 0.29). AMORTIX-PPO trails on all three, most clearly on MPO (0.552 ± 0.015 versus 0.619 ± 0.004).

• Advantage-signal stability. On Kinase MPO, AMORTIX-GRPO keeps the per-epoch standard deviation of the advantage signal around σ ≈ 0.025 throughout training, while AMORTIX-REINFORCE and AMORTIX-PPO plateau one to two orders of magnitude higher — a global baseline or value critic fails to absorb heterogeneous scaffold difficulty.

• PMO benchmark. Running AMORTIX as an instance optimizer under the 10,000-oracle-call budget, it achieves a cumulative top-10 AUC of 16.433 across 22 PMO tasks — the best among amortized policies (f-RAG 16.301, SynGBO 16.426, REINVENT 15.003) and second overall, behind the discrete diffusion instance optimizer GenMol at 17.540.

• Favorable compute crossover. AMORTIX's fixed training budget of approximately 50,000 oracle calls is recouped after roughly 100 test scaffolds against Mol GA and roughly 150 against GenMol on Kinase MPO; at 500 scaffolds, instance optimizers consume on average orders of magnitude more total oracle calls.

• Prodrug transfer. Fine-tuned on only four exemplar parents (morphine, GABA, nipecotic acid, aspirin) and evaluated on 52 held-out drugs with 32 sampled candidates each, the policy recovers three distinct cleavable motifs (alkyl ester, alkyl carbamate, benzyl ester) from a scalar purity reward. 16 of 52 (31%) generations satisfy the paper's definition of a pure prodrug and 33 (63%) carry at least one cleavable handle, with a mean predicted BBB-penetration lift of +0.26 over the parent. On cysteine, the thiol is left unmodified, suggesting the policy learned which groups warrant protection rather than applying a fixed pattern.

Methodology in Plain English

The authors treat molecule construction as a sequential decision process. Starting from a fixed molecular subgraph (or a single atom, which recovers de-novo design), a decoder-only Graph Transformer adds atoms and bonds step by step. Each step is a three-level composite action: choosing an operation (add atom, modify existing, or stop), choosing where to connect it, and specifying the bond order. A validity mask enforces valence rules at every step, so every generated structure is chemically valid by construction. Setting the starting subgraph to one atom also covers de-novo generation.

Before RL, the model is pre-trained on ChEMBL version 35 with supervised teacher forcing so the policy produces plausible drug-like molecules rather than merely valid graphs.

The training innovation is in how rewards are turned into learning signals. Because some scaffolds are easy to optimize and others are chemically constrained, reward magnitudes are not comparable across structures. Instead of a global baseline or a learned value critic, the authors sample G completions from each of B starting structures per training step and compute each completion's advantage as its reward minus the mean reward of its own group. Two departures from the original GRPO objective are deliberate: the division by group standard deviation is dropped (so structures where all completions score similarly do not get disproportionate gradient weight), and normalization by trajectory length is dropped (so concise completions that reach a valid molecule in fewer steps are not penalized). Because sampling trajectories is cheap relative to a parameter update, only a single gradient step is performed per batch, making PPO-style clipping unnecessary.

Evaluation spans three settings: the PMO benchmark under its strict 10,000-call budget (run as an instance optimizer, since each PMO task optimizes one objective from one starting point and the heterogeneity motivating GRPO does not arise); goal-directed scaffold decoration using ZINC-250k scaffolds, Murcko scaffold decomposition, and Butina clustering on Morgan fingerprints at Tanimoto cutoff 0.4, with entire clusters assigned to train/validation/test and 500 scaffolds reserved for testing across three random seeds; and the few-shot prodrug case study.

Why This Matters

Impact on research. The paper reframes a practical obstacle in amortized molecular optimization as a statistical one: reward signals from different starting structures are not on comparable scales, so global-baseline RL methods are dominated by easy instances. Demonstrating that group-relative normalization fixes this gives the field a concrete, critic-free recipe for training conditional generative policies on heterogeneous problems, and provides a clean argument for when amortization beats per-instance search.

Real-world applications:

• High-throughput virtual screening, where an amortized policy can propose optimized candidates for thousands of scaffolds without paying oracle costs per scaffold.
• Lead optimization and scaffold decoration in medicinal chemistry, where a fixed core must be preserved while properties are tuned.
• Prodrug design for blood-brain-barrier penetration, as illustrated by the transfer of a cleavable-handle rule from four exemplar parents to 52 unseen drugs.
• Workflows using expensive oracles such as free-energy perturbation, where the per-call cost makes repeated search per structure prohibitive.

Industry relevance. The economics are the point: AMORTIX front-loads roughly 50,000 oracle calls into one-time training and then answers with zero inference-time oracle calls, whereas Mol GA and GenMol each receive up to 10,000 calls per scaffold in the comparison — a single test case consuming 20% of AMORTIX's entire training budget. For organizations running large structure libraries against costly scoring functions, the break-even point occurs after roughly 100 to 150 test scaffolds, after which the amortized savings grow linearly.

Future Directions

• Beyond 2D topologies. The framework currently handles only 2D molecular topology and does not capture 3D conformational dynamics or stereochemistry, which matter for binding affinity.
• Disconnected starting structures. The action masking assumes a single connected starting scaffold; fragment linking would require modifying the structural constraint logic.
• Richer edit operations. Generation is restricted to additive actions; supporting atom deletions or substitutions would enable more flexible core modifications and scaffold hopping.
• Quantifying the amortization trade-off. The authors state that further experimental validation is needed to identify the specific scale and use cases where the upfront training cost is most justified, and they see potential in hybrid approaches that initialize instance-optimization search with near-optimal candidates from an amortized policy.

Target Audience

Machine learning researchers working on generative models for molecules and RL for structured generation will find the GRPO adaptation and ablation study most directly useful. Computational chemists and drug discovery practitioners evaluating amortized versus per-instance optimization for scaffold decoration, lead optimization, or prodrug design will benefit from the benchmark comparisons and the cost-versus-performance analysis. Practitioners deciding whether to invest in training a reusable policy should note the paper's own caveat that training remains oracle-intensive and the amortization trade-off is not yet fully quantified.

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