MindSpore AI科学计算系列 | 生成扩散模型DiffDock显著提升分子对接准确率

背景

关于药物发现中的分子对接（molecular docking）任务，在此前的一篇专栏文章当中做过一些基本的介绍（https://zhuanlan.zhihu.com/p/492204441）。分子对接过去长期被认为是一个搜索任务，例如分子之间的空间识别常采用格点计算、片断生长等方法，而能量计算则使用模拟退火、遗传算法等方法。随着深度学习的兴起，分子对接被更多的视为一种回归任务，这一变化已经带来了分子对接速度的显著提升，但是结果的准确性仍有待改善。

本次给大家分享的是来自MIT CSAIL的Regina Barzilay教授和Tommi Jaakkola教授课题组的工作，他们将分子对接视为一种生成任务，并采用了时下在图像生成等领域相当热门的生成扩散模型（DGM）。他们的这一模型DiffDock在已知的对接任务中取得了38%的成功率，对比最先进的深度学习模型（20%）和基于搜索的方法（23%）成功率有较大提升，同时速度也有3-12倍提升。对于被折叠的复合物结构，此前各种方法最大仅达到10.4%的对接成功率，而DiffDock仍达到了21.7%的成功率水平。

论文：

DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking

链接：

arxiv.org/abs/2210.01776

代码：

github.com/gcorso/DiffDock/

1、模型

本文的模型概览如图1所示，该模型以单独的配体和蛋白质作为输入结构。随机采样的初始对接位姿在平移、旋转和扭转自由度上反向扩散去噪。被采样的对接位姿按信心模型，产生一系列预测和信心评分，并按照评分排序，获得最终预测。

图1. DiffDock模型概览

1.1 配体位姿

配体位姿（ligand pose）就是配体在三维空间中所有原子的位置，原则上可以认为位姿x是3n维流形（manifold）空间中的一个点，其中n是原子的数量。

然而，对于分子对接来说，自由度远没有3n这么多，因为键长、键角和配体中的环基本上都是刚性的，使得配体的柔性几乎只体现在可旋转键的扭转角上。传统对接方法以及大多数机器学习方法，将3n维流形空间中单独配体的种子构象c作为输入，并仅改变最终结合构型中的相对位置和扭转自由度。因此符合c构象的位姿空间就是一个(m+6)维的子流形空间（m是可旋转键的数量，6来自于对目标对接蛋白质整体的旋转平移）。

关于流形空间的另一种表述是：“任何与种子构象一致的配体姿态都可以通过平移、旋转和扭转角的变化组合达到。作者遵循这个范式，将种子构象c作为输入，并将分子对接作为学习以蛋白质结构y为条件的p_c(x | y)条件概率分布的任务。另外，作者定义了从一个流形到另一个“更好的”流形的一对一映射（变换），其中的扩散kernel可以被直接采样。

1.2 扩散模型

对于配体位姿的以上三种转换（平移、旋转、扭转角变化），随机扩散方程可以被统一定义为如下形式：

其中x为位姿，w为布朗位移，σ^2为噪声的方差。平移情况下的扩散最简单，就是一个标准高斯分布，而旋转和扭转角变化情况下的扩散则略微复杂，具体的实现见Nikolayev & Savyolov, 1970[1]和Leach et al., 2022[2]。

尽管作者将扩散kernel定义在了(m+6)维的子流形空间上，但扩散kernel训练和推理的过程仍在三维坐标空间中进行。这是因为给打分模型提供完整的分子三维结构（而不是流形空间中抽象的一个点），有助于让模型理解分子间的物理相互作用，并且不会被扭转角定义的人为因素所影响，最终也有利于推广到从未见过的复合体的情形。

1.3 信心模型

为获得信心模型d(x, y)的训练数据，作者先运行训练好的扩散模型，为每个训练的样例获得一系列的候选位姿，并生成检测候选位姿的均方根偏差（RMSD）是否小于2Å的T/F标签。然后信心模型就用交叉熵损失来训练，目标是为每个位姿预测一个准确的T/F标签。在推理过程中，扩散模型并行运行并产生N个位姿，并传递给信心模型，最后信心模型会根据获得的RMSD小于2Å的信心值给这些位姿排序。

2、结果

作者使用了PDBBind（一个从PDB库采集的蛋白-配体复合物结构的数据集）上的复合物来评估DiffDock生成10个样本和生成40个样本的效果。所得结果和最佳的基于搜索的分子对接方法SMINA，QuickVina-W，GLIDE，GNINA以及更老的Autodock Vina进行了对比，也和最近的一些基于深度学习的分子对接方法如EquiBind和TANKBind等对比。结果如下：

图2. 基于PDBBind数据的对接结果，与其他分子对接方法的对比

同时作者也量化了DiffDock的表现随着生成样本数的变化，以及选择最佳对接结果的准确性。如下图所示：

图3. DiffDock的表现随生成样本数的变化(左)&选择最佳对接结果的准确性(右)

3、小结和感想

可以看到本项工作的贡献主要有以下几项：

1、将分子对接任务定义为一个生成问题。

2、根据分子对接的自由度，给配体位姿构造了一个新的扩散过程。

3、在PDBBind上实现了38%的Top1预测，大大超过了以前的最佳搜索算法（23%）和深度学习方法（20%）。

总的来说，本工作中将分子对接视为生成式问题的视角是非常新颖的，但也要看到，尽管DiffDock相比过去的一些方法有了准确度的提升，但仍然无法达到完全替代传统方案的水平。期待这一领域有更多采用AI方法的新思路，进一步推动药物发现工具的发展。

参考文献

[1] Dmitry I Nikolayev and Tatjana I Savyolov. Normal distribution on the rotation group so (3). Textures and Microstructures, 29, 1970.

[2] Adam Leach, Sebastian M Schmon, Matteo T Degiacomi, and Chris G Willcocks. Denoising diffusion probabilistic models on so (3) for rotational alignment. In ICLR 2022 Workshop on Geometrical and Topological Representation Learning, 2022