🤖 machine learning

Amortized Molecular Optimization via Group Relative Policy Optimization

The paper introduces AMORTIX, an amortized Graph Transformer model that achieves efficient, single-pass molecular optimization without inference-time oracle calls by employing Group Relative Policy Optimization to stabilize training across diverse structural constraints and starting structures.

Original authors: Muhammad bin Javaid, Hasham Hussain, Ashima Khanna, Berke Kisin, Jonathan Pirnay, Alexander Mitsos, Dominik G. Grimm, Martin Grohe

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Muhammad bin Javaid, Hasham Hussain, Ashima Khanna, Berke Kisin, Jonathan Pirnay, Alexander Mitsos, Dominik G. Grimm, Martin Grohe

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to create the perfect new dish. You have a specific, non-negotiable base ingredient (like a specific type of pasta or a unique spice blend) that must stay in the final recipe. Your goal is to add other ingredients to make it taste amazing, but you have to test every single combination by actually cooking it and tasting it.

In the world of drug discovery, this "tasting" is done by a computer program called an Oracle. It's incredibly expensive and slow to run.

The Old Way: The "Trial and Error" Chef (Instance Optimization)

Currently, most scientists use a method called Instance Optimization. Imagine that for every single new base ingredient you want to improve, you hire a new team of chefs.

They start from scratch.
They cook thousands of variations.
They taste-test them all (calling the expensive Oracle) to find the best one.
Once they find the winner, they throw the whole team away.
If you have 1,000 different base ingredients, you have to hire 1,000 different teams and pay for 1,000 separate, expensive tasting sessions.

This works, but it's slow and costs a fortune if you have many different starting points.

The New Way: The "Super-Learner" Chef (Amortized Optimization)

The paper introduces AMORTIX, a new approach called Amortized Optimization. Think of this as hiring one genius chef who goes to "culinary school" (training) for a while.

The Training: The chef studies thousands of examples, learning general rules like "if the base is spicy, add something sweet" or "if the base is heavy, add something light."
The Payoff: Once the chef is trained, you can hand them any new base ingredient, and they can instantly design the perfect dish in a single second. They don't need to cook and taste thousands of options again; they just use what they learned.
The Benefit: The cost is paid upfront during training. After that, creating new recipes is almost free and instant.

The Big Problem: The "Hard vs. Easy" Puzzle

The authors discovered a major snag with previous "Super-Learner" chefs.

Some base ingredients are easy to improve (like adding salt to bland soup). Almost any change makes them better.
Some base ingredients are hard (like trying to make a rock taste good). Even the best changes only make them slightly better.

If you train a chef by comparing all their dishes together, the "easy" dishes (which get high scores easily) drown out the "hard" dishes. The chef thinks, "Well, I made a great soup, so I'm doing great!" and ignores the fact that they failed to improve the rock. The learning signal gets confused because the difficulty varies so wildly.

The Solution: Group Relative Policy Optimization (GRPO)

AMORTIX solves this with a clever trick called Group Relative Policy Optimization.

Instead of comparing the chef's "easy soup" to their "hard rock" directly, the system puts them in separate groups:

Group A (Easy Bases): The chef makes 10 variations of the easy soup. The system compares them only to each other. "Which of these 10 soups is the best?"
Group B (Hard Bases): The chef makes 10 variations of the rock. The system compares them only to each other. "Which of these 10 rocks is the least terrible?"

By judging the chef against their own peers in the same "difficulty group," the system learns the right lessons for both easy and hard tasks without getting confused.

What Did They Prove?

The authors tested AMORTIX in three main scenarios:

The Standard Test (PMO Benchmark): They played a standard game of "find the best molecule" against other top methods. AMORTIX was the fastest and most efficient, ranking first among the "Super-Learner" methods and second overall.
The "New Ingredient" Test (Scaffold Decoration): They gave the system brand new base structures it had never seen before. AMORTIX successfully improved them instantly, beating other methods that had to re-cook thousands of times for every single new base.
The "Few-Shot" Test (Prodrug Design): They showed the system just four examples of how to turn a drug into a "prodrug" (a drug that activates inside the body). The system learned the rule and successfully applied it to 52 completely different drugs it had never seen, creating valid, working designs instantly.

The Bottom Line

AMORTIX is a new AI tool that learns to design drugs by studying many examples at once, rather than starting over for every new problem. It uses a smart grouping trick to handle both easy and difficult chemical structures, allowing scientists to generate optimized drug candidates instantly without paying the high cost of running thousands of computer simulations for every single new drug idea.

Note: The paper explicitly states that while the AI designs these molecules, it currently only works with 2D chemical structures (flat drawings) and does not yet account for 3D shapes or stereochemistry, which are important for real-world physical drug binding.

Technical Summary: Amortized Molecular Optimization via Group Relative Policy Optimization (AMORTIX)

1. Problem Statement

In drug discovery, molecular design often requires modifying a fixed starting structure (e.g., scaffold decoration, linker design) rather than generating molecules from scratch. These tasks are typically evaluated using expensive "oracles" (scoring functions), ranging from heuristic proxies to high-fidelity simulations like free-energy perturbation.

Current state-of-the-art methods predominantly follow an instance optimization paradigm. Algorithms such as Genetic Algorithms (Mol GA) and guided discrete diffusion models (GenMol) perform an iterative, search-based optimization for every new input structure. This approach requires thousands of oracle calls per design task, resulting in high computational costs that scale linearly with the number of starting structures.

While amortized optimization—learning a transferable policy to generate optimized molecules in a single forward pass without inference-time oracle calls—offers a theoretical solution, existing methods struggle in practice. They fail to generalize to diverse structural constraints or suffer from training instability. The authors identify the root cause as heterogeneity in optimization difficulty: some starting structures yield high-scoring derivatives easily, while others are chemically constrained and yield modest scores regardless of modification. Standard reinforcement learning (RL) methods using global baselines (e.g., REINFORCE) or learned value critics (e.g., PPO) fail to stabilize training under this variance, as the reward signals become incomparable across different starting structures.

2. Methodology: AMORTIX

The authors propose AMORTIX, an amortized Graph Transformer model designed to optimize molecular structures under strict structural constraints with zero inference-time oracle calls.

2.1 Policy Architecture

AMORTIX employs a decoder-only Graph Transformer that operates directly on molecular graph representations. The generation process is modeled as a sequential Markov Decision Process (MDP) with a hierarchical action space:

Operation Selection: Choose to add a new atom, modify an existing atom, or terminate.
Target Selection: Select the specific atom in the current graph to connect to.
Bond Specification: Determine the bond order (single, double, triple, etc.).

To ensure chemical validity, a validity mask is applied at every step, zeroing out actions that would violate atomic valence rules. This guarantees that all generated structures are chemically valid by construction. The model is pre-trained on the ChEMBL database via supervised teacher forcing to learn chemical plausibility before RL fine-tuning.

2.2 Training via Group Relative Policy Optimization (GRPO)

The core innovation lies in adapting Group Relative Policy Optimization (GRPO) to the molecular domain to address the heterogeneity of starting structures.

Group Formation: For each training step, a batch of starting structures is sampled. For each starting structure $S_i$ , the policy generates a group of $G$ completed trajectories (molecules).
Relative Advantage: Instead of using a global baseline or a value critic, the advantage $A_{i,j}$ for a trajectory $j$ from structure $i$ is calculated by centering its reward $R(O_{i,j})$ against the mean reward $\mu_i$ of its specific group:
$A_{i,j} = R(O_{i,j}) - \mu_i, \quad \text{where } \mu_i = \frac{1}{G} \sum_{k=1}^{G} R(O_{i,k})$
Normalization: The authors omit the division by group standard deviation and trajectory length normalization. This prevents structures where all completions score similarly (either trivially easy or near-impossible) from receiving disproportionate gradient weight and avoids penalizing concise completions.
Update: The policy is updated using a single gradient step per batch, aggregating contributions across all generation steps. This approach decouples the learning signal from cross-structure difficulty variations.

3. Key Contributions

Identification of Failure Modes: The paper identifies that high variance in reward signals caused by heterogeneous molecular structure difficulty is a primary failure mode for amortized RL in structurally constrained optimization. Ablation studies demonstrate that neither REINFORCE with a global baseline nor PPO with a value critic effectively stabilizes training in this setting.
AMORTIX Framework: A Graph Transformer-based policy trained with group-relative RL that generalizes to out-of-distribution scaffolds with zero inference-time oracle calls.
Performance Superiority: AMORTIX outperforms both amortized and instance-optimization baselines on goal-directed scaffold decoration tasks while requiring significantly fewer total oracle calls when scaling to multiple starting structures.

4. Experimental Results

The authors evaluate AMORTIX on three distinct tasks:

4.1 PMO Benchmark (Sample Efficiency)

On the PMO benchmark (22 tasks with a 10,000 oracle call budget), AMORTIX was run as an instance optimizer (using standard REINFORCE with a global baseline, as starting structure heterogeneity is absent in single-task de-novo generation).

Result: AMORTIX achieved a cumulative top-10 AUC score of 16.433, ranking first among amortized policies and second overall across all paradigms, trailing only the discrete diffusion instance optimizer GenMol.

4.2 Goal-Directed Scaffold Decoration

This evaluation tests generalization to unseen scaffolds (out-of-distribution) for single-target (GSK3 $\beta$ , JNK3) and multi-target (Kinase MPO) optimization.

Setup: Scaffolds were split by Murcko decomposition and clustering to ensure test scaffolds were topologically distinct from training data.
Results: AMORTIX achieved the highest scores on all three tasks, outperforming amortized baselines (LibINVENT, GraphXForm, DrugEx v3) and instance optimizers (Mol GA, GenMol).
- Kinase MPO: AMORTIX scored 0.619, significantly outperforming LibINVENT (0.455) and instance optimizers Mol GA (0.461) and GenMol (0.479).
- Stability: Ablation studies confirmed that AMORTIX-REINFORCE and AMORTIX-PPO exhibited severe instability (high standard deviation in advantage estimates) compared to the stable AMORTIX-GRPO.
Cost Efficiency: AMORTIX's fixed training budget (~50,000 oracle calls) is recouped after processing roughly 100 test scaffolds against Mol GA and 150 against GenMol. Beyond this break-even point, AMORTIX offers orders of magnitude lower cumulative oracle costs.

4.3 Case Study: Prodrug Transfer

The model was fine-tuned on just four exemplar drugs to learn a modification rule for improving CNS penetration via prodrug design.

Result: Applied to 52 unseen parent drugs, AMORTIX successfully generated structurally clean prodrugs. It recovered three distinct cleavable motifs (alkyl ester, alkyl carbamate, benzyl ester) without explicit instruction.
Performance: 31% of generations were "pure" prodrugs (cleavable junctions only), and 63% contained at least one cleavable handle, with a mean predicted blood-brain barrier (BBB) penetration lift of +0.26.

5. Significance and Claims

The paper claims that AMORTIX demonstrates that the computational burden of molecular optimization can be effectively shifted from inference-time search to training-time learning. By adapting GRPO to normalize rewards within groups of completions sharing the same starting structure, the method stabilizes learning in the presence of heterogeneous difficulty.

The authors position AMORTIX as a scalable alternative to iterative search methods, enabling high-throughput design without the prohibitive cost of repeated oracle evaluations. They note that while the training process remains oracle-intensive, the amortized nature of the policy makes it highly efficient for scenarios involving many starting structures.

6. Limitations

The authors explicitly acknowledge several limitations:

2D Topology: The framework operates exclusively on 2D molecular topologies, not explicitly capturing 3D conformational dynamics or stereochemistry.
Connectivity: The current action masking assumes a single, connected starting scaffold. Extending to disconnected subgraphs (e.g., fragment linking) requires architectural modifications.
Additive Actions: The generation process is restricted to additive actions (adding atoms/bonds). It does not support atom deletions or substitutions, which are necessary for tasks like scaffold hopping.
Training Cost: While inference is cheap, the training process requires tens of thousands of oracle calls, necessitating further validation to precisely quantify the amortization trade-off for specific use cases.