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

TL;DR

The study addresses multi-objective de novo 3D molecular design by coupling diffusion-based generators with an uncertainty-aware reinforcement learning framework. It integrates a conditional EDM backbone, surrogate models for predictive uncertainty, and a PPO-style RL loop to optimize a multi-objective reward $R_{\text{total}}(m)$ that combines $U_{\text{multi}}$ with a reward bonus while penalizing similarity, enabling stable, high-quality generation across multiple datasets. Across QM9, ZINC15, and PubChem, the approach yields higher validity, uniqueness, novelty, and stability than strong baselines, with MD simulations and ADMET profiling supporting drug-like binding potential to EGFR. This framework offers a scalable pathway for automated, uncertainty-guided multi-objective diffusion design with practical implications for early-stage drug discovery.

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.

Figures (9)

• Figure 1: Architecture diagram of the RL-guided diffusion model. (a) Conditional EDM is trained to generate 3D molecules from random atoms, conditioned on target properties. The model learns a forward noising process $q(z_t|z_{t-1}, c)$ and a reverse denoising process $p_\theta(z_{t-1}|z_t, c)$. (b) In the RL phase, the pre-trained conditional EDM is used to generate n molecules by sampling full denoising trajectories. The transition probabilities $p_\theta$ at each timestep are recorded. For each generated 3D molecule, structural diversity is computed, and a reward boost is applied if the molecule satisfies validity, uniqueness, and novelty criteria. In parallel, surrogate models are used to estimate the probability that each molecular property exceeds its cutoff. These probabilities are multiplied to compute the overall success probability (termed “overall uncertainty”, i.e., the likelihood that the molecule satisfies all desired property thresholds.) Rewards are constructed by combining the overall uncertainty, the bonus, and the diversity. These rewards are multiplied along trajectories and used, together with transition log-probabilities, to compute a clipped policy gradient loss for updating the model.
• Figure 2: RL-guided diffusion model training performance across three datasets. (a–c) Evolution of key molecular quality metrics (validity, uniqueness, novelty, atom stability, and molecular stability) over RL training episodes, where the vanilla diffusion models were pre-trained on QM9, ZINC15, and PubChem, respectively. (d–f) Corresponding total reward curves showing convergence trends.
• Figure 3: Distributions of generated valid molecule properties across three datasets with and without RL-guided optimization. (a–c) QED distributions for vanilla diffusion models pre-trained on QM9, ZINC15, and PubChem. (d–f) SAS distributions. (g–i) Binding affinity distributions. Each dataset shows results from three independent trials comparing vanilla diffusion models (without RL) and RL-guided diffusion models (with RL).
• Figure 4: Comparison between a known EGFR inhibitor and top molecules generated by RL-guided diffusion models. (a) Docking pose of a known EGFR inhibitor. (b–d) Docking poses of the top-ranked molecules generated by RL-guided diffusion models pre-trained on QM9, ZINC15, and PubChem, respectively. QED, SAS, and binding affinity values are shown in each panel. (e–l) Corresponding ADMET radar plots and MD simulation RMSD curves. (e, f) Known inhibitor. (g, h) Molecule from (c). (i, j) Molecule from (b). (k, l) Molecule from (d). All MD trajectories were verified to have reached equilibrium prior to analysis, as indicated by the stabilization of RMSD values within a stable range of approximately 0.20–0.30 nm after an initial relaxation phase and the absence of sustained drift or abrupt fluctuations. The presented molecules were selected based on their overall performance.
• Figure B.1: Overall workflow of our framework. A molecular dataset is first preprocessed to train both the diffusion model and the surrogate models. RL guides the optimization of the diffusion model using predictions from the surrogate models. The optimized diffusion model then generates candidate molecules for in silico evaluation.