Nat Commun｜双重扩散模型实现基于靶标口袋的三维分子生成

2024年3月26日，香港城市大学黄家骏老师团队，联合腾讯人工智能实验室和上海锐格医药，在Nature Communications上发表文章A dual diffusion model enables 3D molecule generation and lead optimization based on target pockets。

作者提出了一个基于目标口袋的分子扩散模型（Pocket-based Molecular Diffusion Model，PMDM），用于生成靶向指定靶标的三维分子。PMDM由双重扩散模型，也就是同时具有局部和全局条件的等变扩散模型组成，这使得PMDM能够充分考虑条件概率刻画的蛋白质信息，以有效地生成分子。实验表明模型性能超越了现有方法。

背景

基于结构的生成式化学模型通过探索广阔的化学空间来设计与靶标具有高结合亲和力的配体，在计算机辅助药物发现中至关重要。然而，现有的生成式模型由于自回归采样的局限而面临瓶颈，其计算效率仍有待提升。

方法

图1概述了条件生成模型PMDM，阐明了其结构组件以及涉及训练和抽样的过程。PMDM在正向过程中逐渐引入高斯噪声，同时采用参数化的反向过程迭代消除噪声(图1a)。该模型由两个等变图神经网络组成，分别获得分子嵌入和靶标口袋嵌入(图1b)。为了方便条件生成，作者设计了上下文机制来结合蛋白质口袋的语义和几何信息，利用交叉注意力层来计算分子和蛋白质口袋的注意分数。

此外，PMDM采用双重扩散策略使模型能够识别分子中不同原子间的相互作用力。该策略包括构造两种类型的虚拟边。首先，原子间距离低于局部阈值τl的原子对通过共价局域边成键，因为当原子靠近时，化学键倾向于支配原子间相互作用力。其次，PMDM建立了连接到剩余原子对的全局边，以模拟距离大于局部阈值τl但小于全局阈值τg的原子的范德华力(图1d)。此外，作者还设计了一个符合分子几何系统的平移、旋转、反射和排列等方差的等变动态核。为了保证生成的分子适应结构口袋，在双等变编码器的隐藏状态更新过程中保持口袋位置固定。

在训练阶段，分子及其对应的结合蛋白口袋都被视为三维点云。在PMDM的正向过程中，输入的分子经历了扩散，类似于非平衡态热力学中观察到的现象，采样的时间步长来自于联合分布。同时，蛋白口袋输入作为条件信息保持固定(图1c)。PMDM的主要目标是学习如何逆转这一过程，从而为有条件的数据分布建模。这使得当口袋信息固定时，能够高效地生成具有高结合亲和力的精确分子。在每个时间步，模型输出分数，该分数表示数据点的对数密度。证据下界（ELBO）目标从这些分数中得出，并作为损失函数。

在采样阶段，通过概率采样来初始化数据状态。在给定靶标口袋蛋白的情况下，通过PMDM的双等变编码器获得转移概率。最终分子是由逐步采样生成的。最后，采用argmax函数对分子的原子类型进行识别，也就是直接选择值最大的原子类型，并由模型输出。

图1 PMDM结构图

扩散模型被表述为两个马尔可夫链：扩散过程和反向过程(即去噪过程)。扩散过程根据方差保持计划迭代地向数据中加入高斯噪声，而反向过程则逐步对数据进行细化，直至消除噪声恢复真实数据。扩散模型的改进目标是通过参数化神经网络学习逆向过程。

扩散过程将实际数据分布逐渐扩散为时间设定为T步的预定义噪声分布。每个时间步长的变换被设定为高斯分布。然后将整个过程表示为一个固定的马尔可夫链，该马尔可夫链以方差调度逐渐向数据添加高斯噪声。可以观察到，如果时间步长足够大，最终分布将更接近于标准高斯分布。

为了从扩散过程得到的扩散数据中恢复真实数据，需要设计反向过程。这个过程也是一个参数可学习的马尔可夫链，可以最大化训练数据的概率。由于直接计算似然函数是棘手的，因而扩散模型采用证据下界(ELBO)来优化。实际上，反向过程的最终目标是学习扩散过程中添加的噪声。

与纯扩散模型不同，PMDM是一种条件扩散模型，由口袋蛋白作为条件概率引导分子生成。因此，需要建立模型得到与口袋蛋白结合的配体分布。通过交叉注意力层获得条件口袋蛋白语义信息，可以有效地融合各种模态。具体来说，PMDM设计了一个双等变扩散模型来学习和生成结合分子的几何形状，来模拟局部化学键图和全局距离图。为了保证配体与蛋白质之间的相对距离，采用等变图神经网络EGNN对整个口袋进行处理，该网络以口袋的几何形状作为条件信息。

如图1b所示，首先采用等变图神经网络SchNet对蛋白质语义信息进行编码。SchNet是一个模拟三维空间中分子量子相互作用的图神经网络。它由连续滤波卷积层组成，用于模拟原子系统并保持不变性，为平衡分子和分子动力学轨迹的基准实现先进的性能。同样，PMDM还使用另一个SchNet将配体原子特征投影到中间表示中，通过交叉注意力机制融合蛋白质语义信息和配体隐藏信息，将注意力矩阵转化为标准正态分布。具体来说，将蛋白质信息作为查询来计算注意力得分。交叉注意层的输出将蛋白质语义信息作为条件上下文。

分子的几何形状对旋转和平移是不变的，在设计扩散模型马尔可夫核时应该考虑到这一性质，实际上，可利用等变马尔可夫核来实现这个期望的性质。由于三维分子生成中的分子几何图形可表示为点云，因此，需要手动为点云构建边，以将它们输入到随后的等变马尔可夫核中。具体而言，将长度小于半径τl的边定义为局部边来模拟共价键，将长度在τl和τg之间的边定义为全局边来捕捉范德华力等远距离信息，如图1d所示。作者将局部半径τl设为3埃，它可以包含几乎所有的化学键，而全局半径τg设为6埃。将独热编码的原子特征和具有局部边和全局边的坐标分别输入到双等变编码器中。概括来说，局部等变编码器通过局部边来模拟分子内的力，如真实的化学键，而全局等变编码器通过全局边来捕捉远距离原子之间的相互作用信息，如范德华力。

除了条件蛋白的语义信息外，还需要考虑条件蛋白的空间信息，以确保生成的配体能够贴合口袋结构而不存在碰撞问题。在这里，结合配体和蛋白质作为完整的口袋作为等变核的输入。因此需要以类似方法，构造输入口袋的局部边和全局边。PMDM通过构建这样的边，建模口袋的形状，并通过图神经网络的消息传递过程，保证配体能够独立地聚合邻域信息。由于将口袋空间信息作为条件，在等变核的每一层更新过程中，保持蛋白质的位置不变。

如图1c所示，PMDM样本来自训练每次迭代的均匀分布。从另一个角度来看，它集成了多个小模型来学习逆向过程。与从头开始从标准高斯分布中采样分子噪声的采样策略不同，给定的片段信息应固定为迭代起始点。在这里，采用掩膜策略从头开始模拟采样过程。在每次迭代过程中，分子片段按照相应的时间步长被扩散过程掩膜，原子类型和坐标的识别与从头开始的采样过程相同。最后，将片段数据与去噪部分结合，得到完整的分子。

结果

作者将PMDM与一些具有代表性的方法进行了比较。如表1所示，采用以下指标对比：(1) Vina分数估计配体与靶袋之间的结合亲和力，这是评估生成的分子如何适合特定的蛋白质口袋的最重要的测量，越低越好；(2) 高亲和力是指在测试集中Vina分数高于基线分子所占的百分比，越高越好；(3) QED通过结合几种理想的分子性质来估计分子的药物相似性，越高越好；(4) SA (synthetic accessibility)表示分子的合成可达性，越高越好；(5)Lipinski系数测量药物符合5条Lipinski规则的数量，越高越好；(6) LogP表示辛醇-水分配系数，如果该分子是良好的候选药物，则该系数应在-0.4~5.6之间，由于大多数情况下LogP不超过1，因此可认为越高越好；(7)多样性表示针对每个口袋生成的分子的不相似程度，可认为越高越好；(8)时间是在所有目标上为每个口袋生成100个样本的平均时间，耗时越短越好。PMDM在8个指标的6个中超越了现有方法。

表1 与其他方法对比

作者设计了消融实验，对比了在去除交叉注意力或去除局部信息时的模型性能。消融模型虽然训练时间有所缩短，但在Vina分数、亲和力和Lipinski分数等关键指标上的性能则普遍降低。结果表明交叉注意力和局部信息对PMDM的重要性。

表2 消融实验

作者还进行了案例分析，选择三个口袋蛋白可视化作为案例分析的代表性样本。如图2所示，对于14GS、2RMA和3AF2这三个靶向口袋蛋白，AR-SBDD和DiffSBDD倾向于产生三原子环，而作者提出的模型PMDM避免了这种不稳定环的产生。虽然数据集只包含3%的三原子环，但AR-SBDD方法产生的这种不稳定结构更多，这意味着这些方法陷入局部最优，无法很好地学习数据分布。相反，PMDM可以考虑口袋孔的形状，并产生更大更复杂的环，其Vina分数更低，性能更好。

图2 案例分析

总结

本文提出了一种分子生成模型PMDM，该模型通过结合扩散框架，可以在特定靶标蛋白的条件下一次性生成三维小分子配体。PMDM使用双等变编码器来处理分子不同尺度的(全局和局部)信息。为了实现蛋白质条件生成，PMDM采用交叉注意力机制，通过在共享的高维空间中融合蛋白质表示和配体表示来考虑蛋白质的语义信息，并将整个口袋作为跨神经网络层固定蛋白质空间信息的等变核的输入，考虑蛋白质的结构信息。与现有方法相比，PMDM效率更高，也实现了更好的性能。未来，PMDM有望促进针对特定靶标蛋白的药物设计，加速药物研发进程。

参考文献

[1] Huang et al. A dual diffusion model enables 3D molecule generation and lead optimization based on target pockets. Nat Commun. 2024