[Paper Review] Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization

Siyi Gu, Minkai Xu+7 more|arXiv (Cornell University)|Jul 1, 2024

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper proposes AliDiff, a novel framework that aligns pretrained target-aware molecule diffusion models with desired functional properties using exact energy preference optimization (E²PO). By learning user-defined reward functions for binding affinity and structural rationality, AliDiff achieves state-of-the-art binding energies (up to -7.07 Avg. Vina Score) while maintaining strong molecular properties and diversity on the CrossDocked2020 benchmark.

ABSTRACT

Generating ligand molecules for specific protein targets, known as structure-based drug design, is a fundamental problem in therapeutics development and biological discovery. Recently, target-aware generative models, especially diffusion models, have shown great promise in modeling protein-ligand interactions and generating candidate drugs. However, existing models primarily focus on learning the chemical distribution of all drug candidates, which lacks effective steerability on the chemical quality of model generations. In this paper, we propose a novel and general alignment framework to align pretrained target diffusion models with preferred functional properties, named AliDiff. AliDiff shifts the target-conditioned chemical distribution towards regions with higher binding affinity and structural rationality, specified by user-defined reward functions, via the preference optimization approach. To avoid the overfitting problem in common preference optimization objectives, we further develop an improved Exact Energy Preference Optimization method to yield an exact and efficient alignment of the diffusion models, and provide the closed-form expression for the converged distribution. Empirical studies on the CrossDocked2020 benchmark show that AliDiff can generate molecules with state-of-the-art binding energies with up to -7.07 Avg. Vina Score, while maintaining strong molecular properties. Code is available at https://github.com/MinkaiXu/AliDiff.

Motivation & Objective

• To address the lack of steerability in existing target-aware diffusion models, which generate molecules without prioritizing high binding affinity or desirable molecular properties.
• To develop a general alignment framework that shifts the target-conditioned chemical distribution toward regions of lower binding energy and higher structural rationality.
• To overcome overfitting in preference optimization by introducing an exact energy optimization method with closed-form convergence.
• To maintain strong molecular properties such as drug-likeness (QED), synthetic accessibility (SA), and diversity during high-affinity molecule generation.
• To demonstrate general applicability across different diffusion-based structure-based drug design models.

Proposed method

• The framework uses a unified variational lower bound to align likelihoods of both discrete chemical types and continuous 3D coordinates in diffusion models.
• It introduces Exact Energy Preference Optimization (E²PO), which analytically computes the converged distribution under user-defined reward functions, avoiding overfitting in standard preference optimization.
• E²PO provides a closed-form expression for the updated conditional distribution, enabling exact and efficient alignment of the diffusion model.
• The method integrates reward functions that encode binding energy and structural rationality, guiding the model toward higher-affinity, more stable ligands.
• The approach is applied via fine-tuning on pretrained models such as IPDiff and TargetDiff, using preference signals from ranked molecule pairs.
• Hyperparameter β controls the trade-off between binding affinity and molecular properties, enabling customizable generation for different drug development needs.

Figure 1: High-level illustration of AliDiff . For a protein target, we can have multiple candidate ligands and rank the preference by certain reward functions, e.g . , binding energy. We align the target-aware molecule diffusion model with these preferences by adjusting the conditional likelihoods.

Experimental results

Research questions

• RQ1 Can we effectively steer target-aware molecule diffusion models toward higher binding affinity without sacrificing molecular properties?
• RQ2 How can we avoid overfitting in preference optimization when aligning diffusion models with reward functions?
• RQ3 Does an exact, closed-form solution for energy-based preference optimization lead to better performance than standard methods?
• RQ4 Can the proposed alignment framework be generalized across different diffusion-based structure-based drug design models?
• RQ5 What is the trade-off between binding affinity and molecular properties, and can it be controlled via a single hyperparameter?

Key findings

• AliDiff achieves a state-of-the-art average Vina score of -7.07 on the CrossDocked2020 benchmark, significantly outperforming baseline models.
• The E²PO variant of AliDiff (AliDiff-E²PO) achieves superior binding affinity metrics compared to the direct preference optimization baseline (AliDiff-DPO), while maintaining competitive QED, SA, and diversity scores.
• Fine-tuning TargetDiff with AliDiff (AliDiff-T) improves average Vina score by 6.2%, QED by 16.6%, SA by 2.8%, and diversity by 2.8% over the original model.
• Ablation studies show that β controls the trade-off between binding affinity and molecular properties, with β=1 achieving a balanced performance.
• Lower β values (e.g., 0.01) significantly boost binding affinity but reduce QED and diversity, while higher β values improve drug-likeness at the cost of slightly lower affinity.
• The framework is generally applicable, as demonstrated by consistent improvements across multiple diffusion models, including IPDiff and TargetDiff.

Figure 2: Overview of AliDiff . This workflow can be summarized as 1) For each protein target (pocket) ${\mathbf{p}}$ in the training set, we retrieve two candidate ligands ${\mathbf{m}}$ ; 2) Label the two ligands as wining sample ${\mathbf{m}}^{w}$ and losing sample ${\mathbf{m}}^{l}$ by desirable