Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization

TL;DR

This paper proposes a novel and general alignment framework to align pretrained target diffusion models with preferred functional properties, named AliDiff, and develops an improved Exact Energy Preference Optimization method to yield an exact and efficient alignment of the diffusion models.

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.

Figures (7)

• 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.
• 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 properties, e.g., binding energies; 3) Calculate the preference optimization objective \ref{['eq:loss-alidiff-exact']} and update the molecule diffusion model $p_\theta$.
• Figure 3: Median Vina energy for different generated molecules (TargetDiff, IPDiff, AliDiff) across 100 testing samples, sorted by the median Vina energy of molecules generated from AliDiff.
• Figure 4: Visualizations of reference molecules and generated ligands for protein pockets (1l3l, 2e24) generated by TargetDiff, IPDiff, and AliDiff. Vina score, QED, and SA are reported below.
• Figure 5: Ablation analysis of AliDiff under different $\beta$. Vina Score, High Affinity, QED, and diversity are reported, where blue lines represent AliDiff-DPO, and orange lines represent AliDiff. The dotted lines represent the baseline IPDiff.

Theorems & Definitions (5)

• Theorem 3.1
• Theorem \ref{theorem-kl}
• Lemma C.1
• Lemma C.2
• proof