[论文解读] 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 Science被引用 3

一句话总结

本文提出AliDiff，一种新颖的框架，通过精确能量偏好优化（E²PO）将预训练的目标感知分子扩散模型与期望的功能性质对齐。通过学习用户定义的结合亲和力与结构合理性奖励函数，AliDiff在CrossDocked2020基准上实现了最先进的结合能（最高达-7.07平均Vina得分），同时保持了强分子性质与多样性。

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.

研究动机与目标

• 为解决现有目标感知扩散模型缺乏可调控性的问题，这些模型生成分子时未优先考虑高结合亲和力或理想的分子性质。
• 开发一种通用对齐框架，将目标条件化的化学分布引导至结合能更低、结构更合理的区域。
• 通过引入具有闭式收敛解的精确能量优化方法，克服偏好优化中的过拟合问题。
• 在高亲和力分子生成过程中，保持强分子性质，如类药物性（QED）、合成可及性（SA）和多样性。
• 展示该框架在不同基于扩散的结构药物设计模型中的通用适用性。

提出的方法

• 该框架使用统一的变分下界，对齐扩散模型中离散化学类型与连续3D坐标的可能性。
• 提出精确能量偏好优化（E²PO），通过解析计算在用户定义奖励函数下的收敛分布，避免标准偏好优化中的过拟合。
• E²PO为更新后的条件分布提供闭式表达，实现扩散模型的精确且高效对齐。
• 该方法整合编码结合能与结构合理性的奖励函数，引导模型生成更高亲和力、更稳定的配体。
• 通过在IPDiff和TargetDiff等预训练模型上进行微调，应用该方法，利用排序分子对提供的偏好信号。
• 超参数β控制结合亲和力与分子性质之间的权衡，支持根据不同药物开发需求实现可定制化生成。

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.

实验结果

研究问题

• RQ1 我们能否在不牺牲分子性质的前提下，有效引导目标感知分子扩散模型向更高结合亲和力方向发展？
• RQ2 在使用奖励函数对齐扩散模型时，如何避免偏好优化中的过拟合？
• RQ3 基于能量的偏好优化是否能通过精确的闭式解获得优于标准方法的性能？
• RQ4 所提出的对齐框架是否可泛化至不同基于扩散的结构药物设计模型？
• RQ5 结合亲和力与分子性质之间的权衡如何？是否可通过单一超参数实现控制？

主要发现

• AliDiff在CrossDocked2020基准上实现了-7.07的最先进平均Vina得分，显著优于基线模型。
• AliDiff-E²PO变体在结合亲和力指标上优于直接偏好优化基线（AliDiff-DPO），同时保持了具有竞争力的QED、SA和多样性得分。
• 使用AliDiff对TargetDiff进行微调（AliDiff-T）使平均Vina得分提升6.2%，QED提升16.6%，SA提升2.8%，多样性提升2.8%。
• 消融研究显示，β控制结合亲和力与分子性质之间的权衡，β=1时实现性能平衡。
• 较低的β值（如0.01）显著提升结合亲和力，但降低QED与多样性；较高的β值则提升类药物性，但亲和力略有下降。
• 该框架具有通用适用性，如在IPDiff和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