Source-linked AI summary
PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints
Boqiao Zhang, Godbless James, Sai Krishna Gottipati, Andrew Fitzgibbon
TL;DR
Molecular improvement methods need to optimize properties while remaining synthesizable and diverse, but PGFS selects reactants indirectly and PGFS+ can collapse inputs onto magnet molecules. PGFS++ uses direct discrete reactant selection with input–output similarity encouragement, improving target properties while preserving output diversity and explicit synthesis routes.
Problem
Molecular optimization in unconstrained chemical space may produce molecules without explicit synthesis routes, while PGFS uses indirect reactant-embedding prediction that limits learning effectiveness.
Method
PGFS++ treats each input as the start of a forward-synthesis trajectory, applying learned reaction templates with compatible building blocks and direct global second-reactant selection plus similarity encouragement.
Results
PGFS++ consistently improves over PGFS and its variants on QED and SEH improvement tasks, with competitive performance against REINVENT4 and SynFlowNet.
Takeaways & Limitations
Input–output similarity bonuses mitigate magnet-molecule collapse while improving diversity, similarity, and route-aware synthesizability with competitive QED.
Abstract
from arXiv · showhide
Improving molecular properties, such as drug-likeness or binding affinity, is a recurring task in early-stage drug discovery. However, molecules optimized in an unconstrained chemical space have limited practical value if they cannot be synthesized. Policy Gradient for Forward Synthesis (PGFS) is a synthesis-aware reinforcement learning method for molecular improvement, but its use of reactant embedding prediction makes reactant selection indirect, which, as we show, limits learning effectiveness. We first develop PGFS+, in which reaction templates and second reactants are represented by trainable embedding lookup tables. Combined with a more effective scoring function and RL algorithm, PGFS+ significantly improves the desired property. However, it exposes a reward-hacking failure mode: a powerful reactant search can map diverse input molecules to the same high-reward magnet molecule, improving the reward while collapsing the output diversity. We therefore introduce PGFS++, a synthesis-aware reinforcement learning framework for input-specific molecular improvement. Given an input molecule, PGFS++ treats it as the start of a forward-synthesis trajectory, applies learned reaction templates with compatible in-stock building blocks, and produces a molecule with improved target properties, an explicit synthesis route, and structural similarity to the input. Experiments on molecular improvement tasks show that PGFS++ improves target properties while preserving high output diversity.
1 Introduction
AI methods can optimize molecular properties, but many do not provide explicit synthesis routes. PGFS addresses this gap by constraining molecular improvement to synthetically accessible routes.
- Generative models can produce chemically valid molecules biased toward desired properties such as QED, but usually lack explicit synthesis routes.QED means Quantitative Estimate of Drug-likeness.
- Only a few methods generate both attractive molecules and clear synthesis routes.
- PGFS is a reinforcement-learning method that constrains molecular search to synthetically accessible routes.
- RL-based synthesis-constrained methods typically learn policies from a predefined set of reaction templates.
A. Illustrative Routes of Generating a Magnet Output Molecule B. Model Performance (QED-Diversity Trade-off)
PGFS+ improves reactant selection by directly scoring compatible candidates, but its reward optimization can collapse diverse inputs onto magnet outputs. PGFS++ adds input–output similarity encouragement to preserve diversity.
- A. Illustrative Routes of Generating a Magnet Output Molecule: PGFS selects second reactants through predicted embeddings and nearest-neighbour retrieval, creating a mismatch between continuous policy actions and discrete reactions.
- A. Illustrative Routes of Generating a Magnet Output Molecule: PGFS+ directly scores all template-compatible second-reactant candidates in a learned reaction space instead of predicting an embedding and retrieving neighbours.It also replaces deterministic k-NN selection with sampling during training.
- A. Illustrative Routes of Generating a Magnet Output Molecule: PGFS+ can map many distinct input molecules to the same high-reward output, creating the magnet output problem and reducing input-specific improvement.This collapse occurs when target-property rewards do not constrain input–output relatedness.
- B. Model Performance (QED-Diversity Trade-off): PGFS++ uses learned embeddings for the first reactant and reaction template to form a query for global second-reactant scoring.
- B. Model Performance (QED-Diversity Trade-off): The reaction runner applies the selected reaction and returns its product as the next trajectory state.
- B. Model Performance (QED-Diversity Trade-off): PGFS++ retains direct global second-reactant selection while adding an input–output similarity bonus to encourage input-specific improvement and preserve diversity.
2 Related Work
Molecular generation methods include genetic algorithms, generative models, rule-based synthesis-aware systems, and reinforcement learning. Synthesis-constrained approaches incorporate chemical feasibility into generation, but PGFS extensions retain related limitations.
- Genetic algorithms recursively generate molecules represented as graphs or strings, while synthetic feasibility is typically addressed indirectly.
- Generative approaches such as VAEs and GANs have been explored for drug discovery but exhibit challenges in synthetic feasibility.
- SYNOPSIS generates molecules from available compounds through rule-based chemical modifications, but its optimization is less flexible and generally less effective than newer methods.
- Synthesis-constrained generative models incorporate chemical feasibility directly and can provide explicit reaction routes for desirable molecules.
- PGFS formulates drug discovery as sequential forward synthesis using valid reactions and commercially available building blocks.
- PGFS-TAC and PGFS-MB extend PGFS, but retain similar limitations to the original framework.
3 Methods
PGFS++ models molecular improvement as synthesis-constrained reaction trajectories, directly selecting compatible second reactants and shaping rewards with input–output similarity. Its similarity bonus is designed to reduce magnet-output collapse by tightening an upper bound on diversity loss.
- PGFS++ framework: Each synthesis step selects a reaction template and, for bimolecular reactions, a compatible second reactant from available building blocks.This hierarchical decomposition produces explicit synthesis routes while reducing the complexity of each decision.
- PGFS limitations: PGFS addresses the large second-reactant space through predicted embeddings and k-nearest-neighbor retrieval.The approach predicts an R2 embedding rather than selecting the discrete reactant directly.
- Direct reactant selection: PGFS+ and PGFS++ directly score template-compatible second-reactant candidates using compatibility masks and replace TD3 with PPO for categorical policies.RDKit substructure matching reduces the candidate space, while learned embeddings form a query for scoring masked R2 embeddings.
- Reward hacking and collapse: 12,415 of 12,689 test episodes (97.8%) produced the same output molecule when PGFS+ used only a target-property reward.This magnet-output collapse maps diverse inputs toward a small number of high-reward products rather than input-specific improvements.
- Similarity shaping: PGFS++ adds an input–output similarity bonus and caps rewards below a similarity threshold to discourage magnet-output collapse.The bonus uses Tanimoto similarity between input and output molecules, with c controlling bonus strength and κ defining the reward cap.
- Diversity analysis: The diversity-drop analysis separates editing effects into non-negative edit variance and input–edit covariance, with sufficiently negative covariance termed inward edit bias.Because the similarity bonus favors larger average input–output similarity, it tightens the derived upper bound on diversity drop.
4 Experiments
Experiments evaluate PGFS++ on synthesis-aware molecular improvement using QED, SEH, diversity, synthetic accessibility, input–output similarity, and route-aware synthesizability. PGFS++ outperformed PGFS variants, matched REINVENT4 comparably, and improved diversity-related metrics over PGFS+ while retaining competitive QED.
- Experimental setup: The evaluation used approximately 118,000 molecules, 102 reaction templates, and 2,000 randomly selected test R1 reactants.The templates included 15 unimolecular and 87 bimolecular reactions, and each model improved every test R1 at test time.
- Metrics: The experiments optimized QED and SEH and evaluated diversity, synthetic accessibility, input–output similarity, and Route-Aware Synthesizability.PGFS++ results were averaged over three random seeds, with variability reported using the mean ± standard deviation of seed-wise medians.
- Metrics: Diversity is one minus the average pairwise Tanimoto similarity between Morgan fingerprints of output molecules.The metric uses unordered molecule pairs from the output set and measures structural variation within that set.
- Metrics: RAS is a route-level synthetic accessibility score in [0, 1], computed by multiplying normalized reaction-feasibility and reactant-accessibility factors across route steps.Because each factor lies in [0, 1], multiplication penalizes longer routes; unimolecular steps assign the reactant-accessibility factor a value of 1.
- Main results: PGFS++ consistently outperformed PGFS and its variants on both QED and SEH improvement tasks, while achieving comparable performance to REINVENT4.Unlike PGFS++, REINVENT4 does not produce explicit synthesis routes for generated molecules.
- Ablation study: Both similarity-bonus variants improved output diversity, input–output similarity, and RAS over PGFS+ while maintaining competitive QED; the multiplicative form better preserved QED.With w = c = 0.5, the two forms achieved similar diversity, 1−SA, input–output similarity, and RAS.
5 Conclusion and Discussion
The conclusion presents PGFS++ as a synthesis-constrained reinforcement-learning framework that improves molecular properties while addressing diversity collapse. It also identifies representation, objective, and hyperparameter limitations for future work.
- Contribution: PGFS++ extends PGFS with a discrete action space, pre-computed T–R2 masks, attention-like global R2 scoring, and input–output similarity encouragement.The framework is designed for synthesis-constrained molecular improvement.
- Results: PGFS++ consistently improved over PGFS and its variants and achieved performance competitive with REINVENT4 and SynFlowNet on QED and SEH tasks.The conclusion summarizes the experimental comparison across both molecular-improvement objectives.
- Limitations: The input–output similarity bonus may be sensitive to the hyperparameter values in Eq. 2.The conclusion notes that similarity can improve while other metrics remain relatively strong, but tuning sensitivity remains.
- Future work: Future work includes testing alternative R1 representations and evaluating objectives beyond QED and SEH, such as docking scores.Suggested representations include RLV2, MACCS, and graph-transformer-based alternatives to Morgan fingerprints.
Supplementary Material for PGFS++: Molecular Property Improvement
The supplementary material section is represented only by the phrase “under Synthesis and Diversity Constraints.”
- The supplied supplementary passage identifies the scope as synthesis and diversity constraints.
S1 Algorithm Details
PGFS++ models synthesis as an on-policy PPO trajectory in which the actor selects chemically feasible reaction templates and compatible second reactants directly. The critic evaluates states, while precomputed masks and PPO updates structure training.
- Training paradigm: At each time step, the current molecule is encoded from its Morgan fingerprint as the MDP state.The actor applies a network to FP(R(1)) to obtain a latent representation.
- Training paradigm: The template head samples a chemically feasible reaction template, with an explicit Stop action ending the trajectory without selecting a second reactant.Tmask removes templates infeasible for the current molecule.
- Training paradigm: For non-Stop actions, a query formed from the state and learned template embedding attends over masked building-block embeddings to sample R(2).The joint discrete action is (T, R(2)).
- Training paradigm: The environment applies ForwardReaction(R(1), T, R(2)) and returns the product and a step reward based on the scored molecular transition.The critic is a state-value network sharing the Morgan-fingerprint input.
- Optimization: Training uses on-policy PPO with rollout collection, generalized advantage estimation, clipped policy loss, value loss, and entropy regularization.The main loop samples fresh starting molecules and updates after each rollout.
- Optimization: Early stopping terminates PPO updates when KL(πold, π) exceeds the target threshold KLtarg.Additional PPO hyperparameters are summarized in Table S1.
S2 Evaluation protocol
Evaluation uses specified reaction-template sets and greedy inference with bounded synthesis trajectories. The supplementary protocol distinguishes general and reversible templates and reports matched QED comparisons.
- Templates: Experiments use both reversible and general reaction-template sets, with reversibility defined by RDKit-valid backward reaction recovery.General templates include reactions without reliable automatic reversion.
- Inference: At test time, models greedily select the highest-probability template and reactant for at most 4 reaction steps, with an explicit Stop action.This inference rule applies across the evaluation protocol.
- Reported comparison: Table S2 reports median QED on the test set under general and reversible templates, using matched PGFS++ QED models as the Ours condition.The table is organized around template-setting comparisons.
S3 Additional experimental results
Additional experiments examine template-setting effects, seed variability, novelty relative to MOSES, and implementation details for PGFS and PGFS++. Results show strong QED under reversible settings, competitive novelty, and variability linked to route choices.
- Template-setting results: 0.856 vs. 0.766 median QED is obtained by SynFlowNet under reversible versus general templates, while PGFS++ remains strong in both settings.The reversible setting favors SynFlowNet because its backward policy depends on template reversibility.
- Template-setting results: PGFS++ is compared with REINVENT4 for general templates and SynFlowNet for reversible templates to respect each baseline’s modeling assumptions.SynFlowNet is omitted from the main general-template comparison and retained for reversible templates.
- Seed variability: Target-property results are highly consistent across three random seeds, with molecule-level metrics averaged across seeds before reporting mean ± standard deviation.Table S3 also reports each seed’s median and the aggregate mean and sample standard deviation.
- Seed variability: 2.72 vs. 1.16 average reactions for seeds 654321 and 123456 coincides with substantially different RAS despite similar 1−SA, motivating RAS alongside SA.The difference is attributed mainly to longer routes and different template or R(2) usage.
- Novelty results: PGFS++ attains higher novelty than REINVENT4 at comparable target property and remains competitive with SynFlowNet on novelty and target property.Novelty is computed from maximum Morgan-fingerprint Tanimoto similarity to approximately 1.5 million MOSES molecules.
- Algorithm comparison: The PGFS++ supplementary algorithm samples discrete template and R(2) actions, applies forward reactions, stores rollouts, and updates the policy with PPO.The PGFS algorithm instead uses continuous-action R(2) prediction and nearest-neighbour retrieval in its displayed procedure.