Title: Uncertainty-Aware Multi-Objective Reinforcement Learning-Guided Diffusion Models for 3D De Novo Molecular Design
Venue: NeurIPS 2025
Accepted: yes
Keywords: 3D De Novo Molecular Design

Diffusion Model

Multi-Objective Reinforcement Learning

Uncertainty Quantification

Deep Learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a reinforcement learning framework integrating uncertainty for de novo molecular design diffusion models. The proposed method leverages predictive uncertainty from surrogate models to construct multi-objective reward functions, enabling the generation of drug-like molecules that balance key properties such as drug-likeness, synthetic accessibility, and binding affinity. Evaluations across benchmark datasets and multiple multi-property aggregation methods show performance gains over existing RL baselines. The approach is further validated through molecular dynamics simulations and ADMET profiling.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n-\tPaper is clear and well written.\n-\tSimple method outperforming baselines.\n-\tAuthors extensively compare both against other RL baselines and multi-property optimization methods.\n\nWeaknesses:\n\n-\tRelevance and presentation of MD results (addressed during revisions).\n-\tAuthors do not compare against stablished methods in the molecular domain (addressed during revisions)\n\nQUESTIONS:\n- What are the implications of assuming independence in Eq. 3? How reasonable is this assumpton when a big number of properties are considered in multi-objective optimization?\n- How much influence do the authors think the way they aggregate the multi-objective task to a single-objective in the baselines has on the results?\n- Why do the authors think the pre-trained models on ZINC have so low permance? Why is still possible to rescue these models with Reinforcement Learning and reach quite good performance?\n\nLIMITATIONS:\n-\tThe MD simulation time parameters are poorly reported as steps or frames. Reporting total physical simulation time would be much more informative. It's unclear right now, if the reopted simuations are equilibrated.\n-\tImprovements in already good docking score do not necessarily correlate with better affinity. Confirming high affinity with a more reliable affinity prediction method (such as Free Energy Perturbation) on the 4 molecules in Figure 4 would strengthen the paper.\n-\tGiven that most of the work on fine-tuning diffusion models for optimized chemical spaces has been done in the context of diffusion guidance, I think it would strengthen the contribution to compare against some these works. Some examples are [1,2,3].\n\n[1] PILOT: equivariant diffusion for pocket-conditioned de novo ligand generation with multi-objective guidance via importance sampling. Julian Cremer, Tuan Le, Frank Noé, Djork-Arné Clevert and Kristof T. Schütt. Chem. Sci., 2024,15, 14954-14967. DOI\thttps://doi.org/10.1039/D4SC03523B\n\n[2] Unified Guidance for Geometry-Conditioned Molecular Generation. Sirine Ayadi, Leon Hetzel, Johanna Sommer, Fabian Theis, Stephan Günnemann. 38th Conference on Neural Information Processing Systems (NeurIPS 2024). https://arxiv.org/abs/2501.02526\n\n[3] Generating 3D small binding molecules using shape-conditioned diffusion models with guidance. Ziqi Chen, Bo Peng, Tianhua Zhai, Daniel Adu-Ampratwum and Xia Ning. Nat Mach Intell 7, 758–770 (2025). https://doi.org/10.1038/s42256-025-01030-w", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an uncertainty-aware RL framework to post-train a molecular diffusion model towards multiple objectives. A surrogate model predicts uncertainty estimation, which is used to reweight reward functions to keep objectives balanced. Experiments on three molecule generation benchmarks as well as MD simulations and ADMET profiling tasks validates the effectiveness of the framework.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: the framework introduces predictive uncertainty into multi-objective RL framework.\n\n\n\nWeakness: \n\n1. Although there are many terms and components in the title, the technical details of most components are not novel, making the title and the paper confusing.\n\n   * Diffusion Models for 3D De Novo Molecular Design: simply a conditional EDM, which has been deployed by hundreds of papers.\n   * Uncertainty-Aware: the predictive uncertainty is obtained by an integration with threshold of predicted mean and variance via D-MPNN, which is simple, black-box and not strict. There are more established methods in uncertainty quantification, such as [1]. \n   * Multi-Objective: the aggregation assumes independency and simply multiplies the single-objective rewards. Again, there are tons of papers in multi-objective optimization.\n   * Reinforcement Learning: seems like the major contribution of the paper. The policy update method resembles PPO, so the only novel part is the reward design. What is the real difference between proposed method and baselines?\n\n2. The comparison between proposed method and baselines is unclear. There are many methods to handle multi-objective RL, but the authors ignore them and the baselines are treated in a single-objective manner with a ridiculous binary reward, which is clearly unfair. Also, please include inference-time reward-guided sampling methods as alternative baselines, such as Sequential Monte Carlo (SMC) and Best-of-N.\n\n3. For experiments, QM9 is a dataset with internal 19 types of molecule properties, including some chemical and thermodynamic quantities. I strongly encourage to include this task (predicting multiple real-valued molecule properties) and compare with conditional molecule generation baselines mentioned in GeoLDM [2].\n   \n\n   \n[1] Huang, K., Jin, Y., Candes, E., & Leskovec, J. (2023). Uncertainty quantification over graph with conformalized graph neural networks. *Advances in Neural Information Processing Systems*, *36*, 26699-26721.\n\n[2] Xu, M., Powers, A. S., Dror, R. O., Ermon, S., & Leskovec, J. (2023, July). Geometric latent diffusion models for 3d molecule generation. In *International Conference on Machine Learning* (pp. 38592-38610). PMLR.\n\nQUESTIONS:\nPlease refer to the weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 1, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work RL is incorporated in the training of guided-diffusion models to facilitate generation of 3D molecules satisfying \nmultiple properties, waiving the need to have differentiable surrogate models of the properties of interest. A conditional diffusion model is first trained on the molecular dataset then further optimized using RL with an uncertainty-based multi-objective reward to encourage the generation of molecules satisfying multiple constraints.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper has nice narrative and experiments to demonstrate the capabilities of the proposed model.\n\nQUESTIONS:\n1) Clarify task directions in line 133. Does it refer to minimization or maximization of a property?\n\n2) In the forward process there is conditioning on C in Figure 1 caption, is this needed? This does not match the description of forward process in Appendix A.1.\n\n3) In RL-guided training, it is explained that property combinations satisfying predefined cutoffs (QED > 0.4, SAS < 8, binding affinity < -4.5) are used as conditions.\nIn generation, how the combinations of property values (QED, SAS, Binding affinity) are chosen for conditioning? \n\n4) Does RL-guidance help only with better conditional generation or it is helping with property optimization. i.e., maximizing/minimizing properties? If the later is not the case, how can your method be altered for this purpose?\n\n5) Based on Figure 3 (WO/RL vs W/RL), for ZINC15 dataset, my guess is that the failure of other methods, in satisfying all three properties (TOP % Metric), lies in failing the QED property which is also reflected in the validity metric (Table 1). How is ZINC15 dataset different from other datasets that could justify this significant difference in performance?\n\n6) (Out of curiosity) How does the RL-guided optimization after pretraining on each dataset, affects the quality of conditional generation for samples used for pretraining?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-10"}

comments: SUMMARY:
The paper presents a methodologically sound pipeline, which systematically integrates three key stages: reward design, data sampling, and optimization using Proximal Policy Optimization (PPO). The authors conduct their analysis within a rigorous experimental setting, ensuring a fair and direct comparison against established baselines.

STRENGTHS AND WEAKNESSES:
Strengths
The paper presents a logically sound and well-structured methodological pipeline.

The experimental evaluation is rigorous, employing fair comparisons against relevant baselines.

A key contribution is the novel integration of an uncertainty metric directly into the reward function, which is a compelling idea.

Weaknesses
The primary weakness is the limited methodological novelty. The proposed framework is largely an assembly of well-established components, which may not meet the innovation threshold for this venue.

The choice to demonstrate the method on an explicit-density model (EDM) instead of a more challenging and widely-used latent diffusion model (like GeoLDM) limits the perceived impact. Applying guidance to latent models is a more difficult and generalizable problem, and the current choice may be seen as targeting a simpler use case. Demonstrating success on a latent model would significantly strengthen the contribution.

QUESTIONS:
1. How is uncertainty computed?

2. Is this method applicable to other architectures, such as GeoLDM?

3. How does the RL-based method compare to training-free guidance [1]?

[1] https://arxiv.org/abs/2209.15408

LIMITATIONS:
Please check weakness.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-21