Uncertainty-Aware Multi-Objective Reinforcement Learning-Guided Diffusion Models for 3D De Novo Molecular Design

AI summary

Overview

Research area: Machine learning for molecular generation — specifically reinforcement learning-guided diffusion models for 3D de novo molecular design, with uncertainty quantification used for multi-objective property optimization.

Technical level: Advanced. The paper builds on equivariant diffusion models, surrogate uncertainty estimation, and policy-gradient RL, and assumes familiarity with Markov-chain diffusion processes, graph neural networks, and therapeutic property prediction.

Scope: A single study proposing an uncertainty-aware multi-objective RL framework that fine-tunes 3D molecular diffusion models, evaluated on three molecular datasets against multiple RL and multi-objective baselines, plus molecular docking, Molecular Dynamics (MD) simulations, and ADMET profiling against known EGFR inhibitors.

What This Paper Is About

Diffusion models can generate realistic 3D molecular structures, but they are poor at hitting several drug-relevant property targets at once — here drug-likeness (QED), synthetic accessibility (SAS), and binding affinity to EGFR. Because these properties are computed by external tools and are not differentiable, standard gradient-based guidance does not apply, and the paper argues that RL is a natural fit for optimizing such black-box, non-differentiable rewards.

The goal is a framework that uses surrogate models with predictive uncertainty to shape a reward signal, so a pre-trained 3D diffusion model can be steered toward molecules that simultaneously satisfy all targets without sacrificing validity, uniqueness, or novelty.

Key Contributions

1. An end-to-end framework combining RL, 3D molecular diffusion models, and uncertainty quantification for multi-objective molecular generation. The authors describe it as the first such integrated framework.

2. A reward function built on uncertainty that combines three auxiliary components: a reward boosting mechanism (triggered by validity, uniqueness, and novelty of a generated molecule), a diversity penalty based on average Tanimoto similarity within a batch, and a dynamic cutoff strategy that updates property thresholds from a moving average of previously generated molecules. The design targets reward sparsity and mode collapse.

3. Extensive benchmarking on three datasets (QM9, ZINC15, PubChem) against four RL-guided diffusion baselines adapted from image generation (SFT-PG, DDPO-SF, DDPO-IS, DPOK) and a vanilla diffusion model, plus ablations against 16 alternative multi-objective strategies across four categories and against ablations of the three reward components. The framework is also applied to different diffusion architectures (Conditional EDM, GeoLDM, GFMDiff trained for comparative analysis).

4. Practical validation of top generated candidates through MD simulations and ADMET profiling benchmarked against known EGFR inhibitors.

Main Findings

• Surrogate models are accurate and well calibrated. R² values are 0.95–0.99 for QED and SAS, and 0.86–0.88 for binding affinity. Uncertainty calibration measured by AUCE remains low across datasets and properties, at 0.02–0.10. Nine surrogate models were trained in total (one per property per dataset).

• RL training converges and improves quality metrics. Validity and molecular stability improve over training episodes, and total reward curves grow steadily and stabilize across all three datasets.

• QM9 results. The method achieves the highest validity (98.17%), VUN score (88.90%), and top molecule percentage (28.33%). Validity improves by over 9% relative to the RL baselines (which sit near 88%, e.g. SFT-PG 88.57%, DDPO-SF 88.65%, DDPO-IS 88.82%, DPOK 88.10%, W/O RL 88.55%).

• ZINC15 results. Near-perfect validity (99.02%) and the best performance across metrics, including a 33.40% top molecule ratio. The paper attributes this to ZINC15's drug-like nature matching the property objectives better than the small organic compounds in QM9.

• PubChem results. 100% uniqueness and novelty are maintained, with the largest margin over baselines, but validity is 16.23%. The authors attribute the lower validity to a mismatch between pre-training on simpler molecules and PubChem's high diversity.

• Property distributions shift in the desired direction. Across three trials, RL guidance consistently increased QED and reduced SAS and binding affinity values (the paper defines binding affinity improvement as lower values). The size of the improvement varied by dataset.

• Ablations favor the full design. The method outperforms scalarization-based (WS, POO, MMM, LSDW), constraint-based (NMD, NMD-WS, CP, PFM), gradient-based (SVD, PCGrad, CAGrad, GradVac), and other uncertainty-based (UCB, EI, MVC, BORE) alternatives on key metrics, with uniqueness slightly decreasing as a trade-off. Removing the diversity penalty causes the largest drop (validity 83.55%, VUN 65.77%, top 25.43%) and removing reward boosting also reduces performance (validity 90.00%, top 25.92%) relative to the full method. Notably, replacing the dynamic cutoff with a static cutoff raises validity slightly to 95.73% and VUN to 90.65% but lowers the top molecule percentage to 24.88%, below the full method's 28.33%.

• Docking, MD, and ADMET results are competitive with known EGFR inhibitors. All evaluated candidates maintained stable protein-ligand complexes with RMSD values within or below reference ranges; two candidates showed particularly strong conformational stability. MD trajectories reached equilibrium with RMSD stabilizing within approximately 0.20–0.30 nm after an initial relaxation phase. ADMET analysis indicated good absorption, low CYP inhibition, and minimal toxicity.

• Evaluation scale. Each model generated 2,000 molecules per run over three independent runs, reported as mean ± 95% confidence interval. The authors note that the discrepancy between training curves and Table 1 stems from the evaluation using 2,000 rather than 128 molecules.

• Cutoffs used for the "Top" metric. Relaxed thresholds of QED > 0.4, SAS < 8, and binding affinity < –4.5.

Methodology in Plain English

The pipeline has three stages.

Stage 1 — pre-train a 3D diffusion model. The backbone is a conditional Equivariant Diffusion Model (EDM) that learns to generate 3D molecules (atomic coordinates and atom features) from noise, conditioned on a target-property vector. The reverse denoising step is parameterized by an E(n)-equivariant Graph Neural Network (EGNN). GeoLDM and GFMDiff are also trained for architecture comparisons.

Stage 2 — train surrogate predictors that know what they don't know. For each property in each dataset, a Chemprop D-MPNN surrogate is trained to output a predictive mean and variance. Assuming a Gaussian prediction, the framework computes the probability that a molecule exceeds a property threshold; for properties where higher is better (e.g., QED) the direction is +1, and for properties where lower is better (e.g., SAS, binding affinity) it is −1. The per-property probabilities are multiplied together into a single joint "overall uncertainty" score in [0, 1] — the estimated likelihood that a molecule satisfies all targets at once.

Stage 3 — fine-tune the diffusion model with RL. The pre-trained model samples full denoising trajectories, recording transition probabilities at each step, rewritten in probability density function form so a likelihood ratio can be computed. Each generated molecule receives a total reward that multiplies the joint satisfaction probability by a bonus for validity, uniqueness, and novelty, then subtracts a diversity penalty (average Tanimoto similarity to other molecules in the batch) weighted by a coefficient that decays exponentially over episodes — encouraging exploration early and exploitation later. Property thresholds are updated dynamically from a moving average of earlier outputs rather than fixed in advance. Parameters are updated with a PPO-style clipped surrogate objective using the ratio between current and previous policies.

Validation. Molecules are scored for validity, uniqueness, novelty, their joint VUN product, atom-level stability, molecular-stability, and the proportion that satisfy all three property constraints. Top candidates are docked (AutoDock Vina) against EGFR, simulated in AmberTools/OpenMM with ff14SB and GAFF parameters in TIP3P water at 300 K and 1 bar for 4,000 ps (approximately 1,000,000 steps), and profiled with ADMET-AI. All experiments ran on NVIDIA A100 GPUs with 80 GB memory. Data was split 80% training, 10% validation, 10% test, with diffusion models split by species grouping and surrogate models split by molecular scaffold.

Why This Matters

The paper's central claim is that uncertainty-aware RL gives a practical handle on multi-objective molecular design where reward functions are black boxes — the usual situation in drug discovery, where properties come from external predictors and docking tools rather than differentiable formulas. It also argues that RL-guided optimization had previously been applied mainly to 1D SMILES strings and 2D graphs, leaving 3D de novo generation underexplored despite being necessary for downstream tasks like docking and MD.

Real-world applications:

• Early-stage drug discovery against EGFR, a receptor protein involved in cancer progression and drug resistance, where candidate molecules must combine drug-likeness, synthesizability, and binding affinity.
• Virtual screening and lead optimization, using the surrogate-plus-RL loop to generate candidates that already pass loose property thresholds before expensive simulation.
• ADMET-aware prioritization, filtering generated molecules for absorption, CYP inhibition, and toxicity before committing to synthesis.
• Beyond pharma, the authors state the approach may benefit materials science, catalyst design, and molecular engineering.

Industry relevance: The framework targets the front end of the drug pipeline, where the cost of exploring chemical space manually is prohibitive. The released implementation (https://github.com/Kyle4490/RL-Diffusion) and the demonstration that generated molecules match or outperform known EGFR inhibitors in docking, MD stability, and ADMET profiling are the pieces most relevant to groups evaluating generative design for real programs.

Future Directions

• Scalable diffusion backbones for large molecules. The PubChem results (validity 16.23%) lead the authors to conclude that current diffusion architectures struggle with large, complex molecules containing many heavy atoms, and that this is a backbone limitation rather than an RL limitation. They explicitly call for future work on scalable architectures.

• Broadening and deepening uncertainty-based multi-objective RL. The paper notes that uncertainty-based methods remain underexplored in RL-guided generative modeling, and its ablation covers only four such alternatives (UCB, EI, MVC, BORE).

• Extending the framework to additional property targets and beyond drug discovery, given the stated potential in materials science, catalyst design, and molecular engineering.

• Experimental validation. The current evidence comes from docking, MD simulation, and ADMET prediction against known EGFR inhibitors; the appendix details on surrogate performance, cross-architecture results, and extended generated-molecule examples are truncated in the provided content and are not reported here.

Target Audience

Researchers and practitioners in generative molecular design, RL for scientific discovery, and computational drug discovery — particularly those working on 3D diffusion models, multi-objective optimization, or uncertainty-aware reward shaping. It is also relevant to medicinal chemists and computational biologists evaluating generative candidates against EGFR, and to machine learning researchers interested in transferring RL fine-tuning techniques from image generation to structured scientific domains. Readers without grounding in diffusion models, equivariant networks, or policy-gradient RL will need to consult the appendices, which the paper references heavily for mathematical derivations and implementation details.

Authors’ abstract

Designing de novo 3D molecules with desirable properties remains a fundamental challenge in drug discovery and molecular engineering. While diffusion models have demonstrated remarkable capabilities in generating high-quality 3D molecular structures, they often struggle to effectively control complex multi-objective constraints critical for real-world applications. In this study, we propose an uncertainty-aware Reinforcement Learning (RL) framework to guide the optimization of 3D molecular diffusion models toward multiple property objectives while enhancing the overall quality of the generated molecules. Our method leverages surrogate models with predictive uncertainty estimation to dynamically shape reward functions, facilitating balance across multiple optimization objectives. We comprehensively evaluate our framework across three benchmark datasets and multiple diffusion model architectures, consistently outperforming baselines for molecular quality and property optimization. Additionally, Molecular Dynamics (MD) simulations and ADMET profiling of top generated candidates indicate promising drug-like behavior and binding stability, comparable to known Epidermal Growth Factor Receptor (EGFR) inhibitors. Our results demonstrate the strong potential of RL-guided generative diffusion models for advancing automated molecular design.