3D EQUIVARIANT DIFFUSION FOR TARGET-AWARE MOLECULE GENERATION AND AFFINITY PREDICTION
论文解读：
3 方法
3.1 问题定义
蛋白质结合位点表示为一组原子
其中 NP 表示蛋白质中的原子数量，在三维空间坐标系中用 xP ∈ R³ 来表示各个原子的位置坐标，在高维特征空间中用 vP ∈ R^Nf 来表示每个原子所具有的元素类型以及氨基酸类型的特征向量。研究的核心目标是基于给定的蛋白质靶点位点信息（protein target），系统地生成与其相互作用作用域相关的结合分子
为了便于描述起见，并非出于其他考虑因素而特意如此设计
博主注：one-hot 编码用于将离散的分类标签转换为二进制向量。
3.2 TARGETDIFF 总览
如前所述，在第1节和第2节中所探讨的，我们希望设计一个非自回归架构以规避由自回归采样方法所带来的局限性。同时，为了更好地模拟蛋白质-配体相互作用的整体特性，在连续的3D空间中进行表征将有助于减少离散化处理带来的误差。此外还需要注意的一点是该系统必须具备全局平移与旋转不变性以保证SE（3）等价性。
鉴于此
的潜在变量，其中 M1 ， M2 ， · ·，MT 是一系列与数据具有相同维数的潜在变量 M0 ~ p（M0 | p )。
如图1所示所述内容中所提及的方法，则包含正向扩散流程与逆向生成流程两个主要组成部分。其中每一个具体流程都被明确地被定义为一个马尔可夫链。其中扩散流程则逐步地将噪声注入到数据中；而生成流程则通过参数化的方式运用θ网络从噪声分布中重构出原始数据分布：
为实现基于给定的蛋白质结合位点生成三维分子的目的, 模型必须生成连续的原子坐标以及离散的原子类型. 同时在整个生成过程中维持SE(3)等距变换特性. 在下一节中, 我们将详细阐述构建扩散过程的方法, 并探讨参数化生成过程的技术细节, 最终完成模型的训练工作.
3.3 分子扩散过程
基于最近关于连续分布和离散分布进展的研究 Ho et al.（2020）与 Hoogeboom et al.（2021），我们采用高斯分布在连续位置信息 x 上以及分类分布在元素种类上进行建模以模拟分子系统的概率结构。为此我们需要定义一种能够反映不同元素特性和芳香环属性的新模型框架。具体而言我们假设分子系统的概率分布由各点位置的概率密度函数 f(x) 和各点所属元素类型的条件概率 P(c|x) 综合决定。在这个假设下我们可以将整个系统的概率密度函数表示为 f(x,c) = f(x)P(c|x) 即各点位置独立于其所属元素类型的联合概率密度函数。为了构建这一模型框架我们首先在每个时间步长 t 中对所有类别进行操作并在此基础上对每个类别中的样本分别施加微小高斯噪声和平移操作以模仿扩散过程中的随机扰动机制。随后通过设计一个基于方差参数序列 β₁, β₂, …, β_T 的马尔可夫链逐步构建这一复杂的系统结构
我们发现，在实际应用中时间表可能会有所变化。然而为了保持简洁我们将联合分子分布分解为扩散过程中的两个独立分布的乘积这是因为这些独立分布具有简单的数学表达式从而能够有效地提取噪声样本。在后续章节中我们将探讨模型如何在生成过程中考虑到原子坐标与类型之间的相互作用
Denoting α_t as 1 − β_t and ᾱ_t as the product from s=1 to t of α_s, a key characteristic of the diffusion process is its ability to determine the noisy data distribution q(M_t | M_0) for any time step, expressed in closed form.
基于贝叶斯理论，在计算原子核位置的正态后验分布和原子类型的分类后验分布时，都可以通过解析方法求解。
3.4 等变性分子生成过程的参数化？
生成机制与其相对应的逆过程会通过从初始噪声MT中恢复真实分子M0来实现。我们通过一个由θ参数化表示的神经网络来模拟这个反向分布。
生成过程所具有的一个性质即是Likelihood在蛋白质-配体复合物经过旋转和平移之后保持恒定,而这即用于生成三维目标（如分子）的关键偏差. 这一结论的重要论据是,基于等变转移函数构建的不变分布将产生一致的结果. 因此,我们在此设定中提出了如下假设:
命题 1. SE(3)变换以Tg的形式进行表征。当质心（CoM）发生位移时，在蛋白质-配体复合物中实现了对Tg的不变似然特性：即当pθ(Tg(M0|P))与pθ(M0|P)相等时（其中蛋白质原子数为零），通过SE(3)-等变网络来参数化马尔可夫转移过程p(xt−1|xt,xP)，从而实现了对动态过程的建模。
A slight misapplication of notation is employed in this text, where xt (t = 1, ..., T) is used to represent ligand atom coordinates and xP for protein atom coordinates. The process of shifting the center of mass is examined in Appendix B and proven invariant likelihood in Appendix C. Multiple parameterization strategies for µθ ([xt , vt ], t, P) and cθ ([xt , vt ], t, P) exist. We opt to utilize neural networks for predicting [x0 , v0 ] and then apply equation 4 to obtain µθ and cθ, which define the posterior distributions. Recent advancements in the field of equivariant neural networks have spurred research efforts: (Thomas et al., 2018; Fuchs et al., 2020; Satorras et al., 2021b; Guan et al., 2022). To model the interaction between a molecule's atoms and a protein's atoms, we employ an SE(3)-equivariant graph neural network.
Within the l-th layer, the atom hidden embedding h and coordinates x are alternately updated in an alternating manner as described below.
在其中d_{ij} = ||x_i - x_j|| 表示原子i与j之间的欧几里得距离，在这种情况下e_{ij} 是一个附加特征，在这种情况下它表示连接的存在性：即连接存在于蛋白质-蛋白质、配体-配体或者配体-蛋白质之间。考虑到我们不希望更新蛋白质基质上的原子坐标，在这种情况下（即masking掉ligand分子），初始隐藏态嵌入h_0是由编码层直接生成的一个嵌入向量来表示所有基质中原子的信息状态。经过一系列多层感知器处理后，在应用了Softmax函数之后得到最终状态向量\hat{v}_L. 由于\hat{x}L对x_t具有旋转不变性（Rotation Equivariance），并且根据方程4我们可以直观地得出\hat{x}{L-1}同样也具有对x_0的旋转不变性性质；从而我们成功地证明了所期望实现的一种马尔可夫变换关系式；完整证明过程详见附录A.
3.5 训练目标
在Kingma和Welling（2013）的研究中，q与p的组合构成了一个变分自编码器。该模型可以通过优化变分下界来实现负对数似然的最小化。对于原子坐标的位置损失而言，在q(x_{t-1}|x_t,x_0)和p_θ(x_{t-1}|x_t)均为高斯分布的情况下，KL散度具有一个闭式表达式：
其中
C被视为一个常数。实际上，在采用无加权的MSE损失（设定γt=1）进行模型训练时也能够获得更好的效果，并非必须依赖加权方法（如Ho et al.（2020）所述）。对于原子类型的损失而言，则是直接计算了基于分类分布的KL散度。
最终总损失由坐标位置损失与类型分类损失的加权总和构成：L = Lt−1 + λLt−1，在附录E中我们对整体训练流程与采样策略进行了详细阐述。
3.6 非监督亲和力排名和预测 （AS UNSUPERVISED LEARNER？）
生成模型可被视为一种无监督学习工具，在蛋白质-配体相互作用研究中展现出巨大潜力。
然而，在目标感知分子生成领域中，
目前尚未明确如何将生成模型与结合亲和力进行有效整合，
这成为评估分子质量的重要指标。
现有方法在估计分子质量方面存在局限性，
尤其是依赖自回归采样的机制，
它们在对数据分布进行建模时通常依赖于数据的自然顺序。
为了克服这些局限性，
我们开发了一种新的参数化架构来关联无监督生成器与亲和力排序关系。
在此架构下，
网络能够预测去噪后的 [x̂0 , v̂0 ]。
对于给定的蛋白质-配体复合物，
我们在冻结 x 更新分支的同时提供 [x0 , v0 ] 给 φθ 网络（即仅更新原子隐藏嵌入 h ）。
经过这一过程后，
最终能够获得 hL 和 v̂0 的表示。
我们假设当配体分子与蛋白质之间存在良好的结合亲和力时，则认为其原子类型应具有较低的灵活性，并且这种关系可以通过计算v̂0的熵来衡量。因此该值可作为评估指标用于排序，并且其有效性已在实验中得到验证。此外hL值还包括一些全局信息。我们发现通过整合这一特性及其相对简单的线性变换处理方法，在实验中显著提升了结合亲和力的排序性能
复现（翻译自github官方项目页）：
Target-Aware Molecule Generation
数据
训练数据/测试数据位于谷歌云盘中的data文件夹中（博主注：需要魔法）。
如果你要从头训练模型，需要下载预处理过的lmdb文件和split文件:
该种类型的数据集版本号为v₁.₁均方根差数量为pocket₁₀最终处理后的.lmdb该种类型的数据集数量为pocket₁₀数据分割成poses.pt
为了在测试集上评估模型性能, 需要下载并解压 test_set.zip. 其中包含了专门用于 Vina.Docking 的原始 PDB 文件.
为了从头开始处理数据集，请访问此处获取CrossDocked 。”然后将此文件包放置于 data/Cross -Docker . 并启动其中包含必要的预处理脚本。
- clean_crossdocked.py 用于 过滤原始数据集，保留那些满足RMSD < 1A的。 它将生成
index.pkl文件，并创建一个包含原始过滤后数据的新目录 ( 与谷歌云盘中的crossdocked_v1.1_rmsd1.0.tar.gz对应)。如果你已经下载了.lmdb文件，你不需要这些文件。
python scripts/data_preparation/clean_crossdocked.py --source data/CrossDocked2020 --dest data/crossdocked_v1.1_rmsd1.0 --rmsd_thr 1.0- extract_pockets.py 会把原始蛋白质文件切割，保留与它结合的分子附近 10A 的区域。E.g.
python scripts/data_preparation/extract_pockets.py --source data/crossdocked_v1.1_rmsd1.0 --dest data/crossdocked_v1.1_rmsd1.0_pocket10- split_pl_dataset.py 用于分割训练和测试集。我们使用与AR 和 Pocket2Mol 相同的分割文件（？）
split_by_name.pt，也可以从谷歌云盘的data文件夹下载。
python scripts/data_preparation/split_pl_dataset.py --path data/crossdocked_v1.1_rmsd1.0_pocket10 --dest data/crossdocked_pocket10_pose_split.pt --fixed_split data/split_by_name.pt训练
从头训练
python scripts/train_diffusion.py configs/training.ymlTrained model checkpoint
请查看详细信息
获取该资源包
采样
从测试集中取样蛋白质口袋
python scripts/sample_diffusion.py configs/sampling.yml --data_id {i} # Replace {i} with the index of the data. i should be between 0 and 99 for the testset.可以使用多GPU加速，例如:
CUDA_VISIBLE_DEVICES=0 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 0 0
CUDA_VISIBLE_DEVICES=1 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 1 0
CUDA_VISIBLE_DEVICES=2 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 2 0
CUDA_VISIBLE_DEVICES=3 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 3 0从pdb文件中采样
为了采自该蛋白质袋进行采样分析，请执行以下代码块
评估
Evaluation from sampling results
The docking mode can be chosen from {qvina, vina_score, vina_dock, none}
Note: Preparing pqdqt and pqr files will require some time when running the evaluation code with vina_score/vina_dock docking mode for the first time.
Evaluation from meta files
We offer the outputs (including docked outputs) of our model and CVAE, AR, and Pocket2Mol as baselines here.
| Metafile Name | Original Paper |
|---|---|
| crossdocked_test_vina_docked.pt | - |
| cvae_vina_docked.pt | liGAN |
| ar_vina_docked.pt | AR |
| pocket2mol_vina_docked.pt | Pocket2Mol |
| targetdiff_vina_docked.pt | TargetDiff |
You can directly evaluate from the meta file, e.g.:
By means of notebooks/summary.ipynb, one can access and review the research findings published in the paper efficiently through this link: notebooks/summary.ipynb.
Binding Affinity Prediction
数据
In the unsupervised learning framework, we continue to utilize the CrossDocked2020 dataset and identify the data with experimentally measured binding affinity (saved in affinity_info.pkl) for further analysis.
Within a supervised learning framework, the PDBBind dataset is utilized. It is accessible via Welcome to PDBbind-CN database. Refined or general versions of the dataset should ideally be saved within the data/pdbbind_v{YEAR} directory.
Take the PDBBind v2016 for example, you need to first unzip the data:
通过以下步骤提取10A-pocket并划分数据集：
Training
One can train the binding affinity prediction model with:
It is also feasible to improve the model by incorporating additional features derived from the unsupervised generative model. You are required to initially export the hidden states using:
This command will output various meta information, and please specify the feature you desire in the subsequent configuration of the supervised prediction model (for example, in configs/prop/pdbbind_general_egnn.yml).
Trained model checkpoint
NOTE: In the context of supervised learning, due to an accidental loss of training results from PDBBind v2020, we currently offer a pre-trained model checkpoint based solely on data from PDBBind v2016 within our initial testing framework. However, despite this limitation, it already achieves sufficiently accurate predictions for real-world applications. We plan to retrain these models using updated data from PDBBind v2020 and distribute new checkpoints as soon as feasible.
https://drive.google.com/drive/folders/1-ftaIrTXjWFhw3-0Twkrs5m0yX6CNarz?usp=share_link
评估
To evaluate unsupervised learning performance, please refer to the notebook at notebooks/analyze_affinity.ipynb.
To perform supervised learning model evaluation, the following command can be utilized for evaluating performance on a test set.
Expected results:
| RMSE | MAE | R^2 | Pearson | Spearman |
|---|---|---|---|---|
| 1.316 | 1.031 | 0.633 | 0.797 | 0.782 |