Nat. Mach. Intell. | CGN+CAN=CGN！Frad大尺度扰动分子提升属性预测

今天为大家介绍的是来自清华大学兰艳艳团队的一篇论文。深度学习方法被认为在加速药物发现和材料设计中的分子筛选方面具有前景。由于标记数据的有限可用性，各种自监督的分子预训练方法相继出现。尽管许多现有方法利用计算机视觉和自然语言处理中的常见预训练任务，但往往忽视了支配分子的基本物理原理。相比之下，在预训练中应用去噪可以视为一种等效的力学习，但有限的噪声分布会引入对分子分布的偏差。为了解决这一问题，作者提出了一种名为分数去噪的分子预训练框架，它将噪声设计与力学习等效所施加的约束解耦。通过这种方式，噪声变得可定制，从而允许结合化学先验，显著改善分子分布建模。实验表明，作者的框架在力预测、量子化学性质和结合亲和力任务上始终优于现有方法，确立了最新的最佳结果。改进的噪声设计提高了力的准确性和采样覆盖率，有助于创建物理一致的分子表征，最终提升预测性能。

分子属性预测在药物发现和材料设计等领域至关重要，但传统方法如第一性原理计算和实验成本高昂，限制了高通量筛选。深度学习被认为可以有效降低成本并加速筛选。然而，缺乏标记分子数据是主要困难，因此提出了各种自监督分子预训练方法，主要借鉴了计算机视觉和自然语言处理的策略。

现有预训练方法往往忽视分子的化学特性和物理原理，例如，图级增强会显著改变分子属性，导致属性预测效果不佳。因此，将化学先验和物理法则融入人工智能是设计合适预训练方法的关键。近期，去噪方法引入了可物理解释的预训练任务，通过学习近似原子力来处理噪声扰动的构象。然而，以坐标高斯噪声（CGN）为噪声类型限制了分子分布的多样性。

作者提出了一种名为分数去噪（Frad）的分子预训练框架，能够有效建模分子分布并保持力学习的物理解释。Frad结合了化学感知噪声（CAN）和CGN，并通过实验验证其在能量面探索、原子力学习和适应多种分子系统方面的优越性。在21个量子化学性质和结合亲和力预测任务中，Frad取得了18项最佳结果。

为了实现物理一致的自监督预训练，作者提出了Frad框架，相当于学习分子中的近似原子力。整个流程如图1所示。

图 1

给定一个平衡分子构象xeq，作者加入了化学感知噪声（CAN）和坐标高斯噪声（CGN）的混合，其中平衡构象指的是分子潜在能量表面的局部极小值结构。然后，模型被训练以从噪声构象中预测CGN，来恢复部分引入的噪声。具体而言，作者使用等变图神经网络从噪声构象中提取特征，然后噪声预测头根据这些特征预测CGN。如图1f所示，作者使用TorchMD-NET作为主干模型，从3D分子输入中获取等变特征。理论分析表明，该任务无论CAN分布如何，均具备力学习的解释，而CAN分布则影响力目标和采样分布。因此，作者精心设计CAN，使其与真实的分子构象分布对齐，从而相较于现有去噪方法获得更精确的力目标和更广泛的采样分布。

作者提出了一个定理，证明了Frad与分子中原子力学习的等价性，从而增强了去噪任务的可解释性。不同于以往的工作，作者寻求该等价性的最小条件，提供了更大的处理噪声分布的灵活性。首先，引入Boltzmann分布假设：构象出现的概率与其能量有关，具体为：

其中，x代表分子构象，E(x)为势能函数，T为温度，k为Boltzmann常数。由此可得，构象分布的得分函数等于分子力的常数倍，即：

去噪预训练数据集包含大量分子的平衡构象，通过添加混合噪声来模拟分子分布。定理证明，在Frad的力学习中，如果混合噪声分布满足CGN分布，则Frad等价于通过Boltzmann分布学习近似的原子力。通过Frad，分子内在规律可以被有效捕捉并推广到多个下游任务。

为了准确模拟分子分布并实现真实的构象采样和精确的力目标，需精心设计混合噪声来捕捉平衡态周围的真实分布。定理1仅对CGN施加了约束，因而CAN有较大的灵活性用于描述分子分布。一种基本的混合噪声设计是仅使用CGN，不引入CAN，通过选择较小的噪声值可以满足分布要求。这就是常见的Coord方法。然而，CGN只能捕捉小尺度的振动，无法涵盖实际分子分布中更广泛的低能量构象。

为了解决这一问题，作者引入CAN以捕捉分子分布的复杂特性。首先，作者提出RN，通过高斯噪声扰动可旋转键的扭转角。然后，为更全面描述各向异性振动，作者提出VRN，使用独立的高斯噪声扰动键长、键角和扭转角，包括可旋转键的扭转角，以更好地模拟分子构象变化。

为了评估Frad在分子属性预测中的有效性，作者进行了系列挑战性下游任务，包括原子级力预测、分子级量子化学属性预测和蛋白质-配体复合物的结合亲和力预测。作者的模型与多个基线进行了系统比较，包括预训练方法和未进行预训练的属性预测模型。在实验结果中，使用缩写Coord指代另一项工作中的Coord预训练方法，该方法与作者的模型共享相同的主干。数据划分方法遵循文献中的标准做法，其中MD17、MD22和QM9采用均匀随机划分，而ISO17和LBA则使用离散分布划分设置。

由于Frad理论上涉及学习近似力，因此预计在下游力学习任务中会带来优势。作者评估Frad在MD17、ISO17和MD22等多个分子动力学数据集上的力预测性能。为了测试Frad在不同分子分布中的推广能力，作者选择MD17和MD22数据集，后者包含更多原子的分子。ISO17数据集包含129种C7O2H10异构体，用于进一步测试Frad对多样化化学结构的适应性。由于许多研究使用MD17作为评估数据集，作者将Frad与几种竞争性的预训练方法和监督模型进行比较。在数据拆分方面，图2展示了在大训练数据和有限训练数据两种场景下的实验结果。

图 2

在这两种情况下，Frad在八种分子中有七种表现优于基线，特别是在与相同主干的去噪预训练方法（如3D-EMGP和Coord）比较时，Frad的优势明显。Frad的预训练在训练集有限时也能显著提升主干模型TorchMD-NET，显示出其有效学习分子的内在原理和良好的泛化能力。

为了验证Frad在不同性质上的竞争力，作者在QM9数据集上评估Frad(RN)和Frad(VRN)的表现。表2显示，Frad在12个任务中超越了大多数监督和预训练方法，并在9个目标上实现了最新成果。此外，Frad在11个目标上相较于TorchMD-NET主干模型有显著提升，表明Frad预训练所学的知识对大多数下游任务有帮助。

在表格下半部分，Frad对去噪预训练方法表现出明显优势，尤其在所有12个任务中，Frad的结果均达到或超过了与其共享主干的Coord方法。这表明引入的化学先验对多个属性类别有所贡献。

蛋白质-配体结合亲和力（LBA）预测旨在预测蛋白质与配体之间的相互作用强度。作者在ATOM3D数据集上评估Frad在LBA任务中的表现，该数据集提供了蛋白质-配体复合物的构象和结合亲和力标签。Frad提取蛋白质的结合口袋，并利用口袋-配体复合结构作为输入预测结合亲和力，结果见表3。

Frad在性能上优于专为蛋白质表示设计的序列和结构基线，表明Frad是一种适用于小分子和蛋白质的通用表示学习方法。此外，Frad也超越了传统的去噪方法，再次凸显了引入化学先验的重要性。

有效的三维分子预训练需要大规模的平衡构象数据集。然而，构建此类数据集通常需要使用密度泛函理论（DFT），成本较高。因此，作者探讨模型对预训练数据准确性的敏感性，评估Frad在使用快速但精度较低的方法计算构象时是否仍然有效。作者使用RDKit在原始PCQM4Mv2预训练数据集上重新生成三维构象，尽管其准确性较低，但速度更快。随后，作者对Frad（τ = 0.04, σ = 2）和Coord（τ = 0.04）进行不准确数据集的预训练。

结果表明，基于不准确构象的预训练导致更大的平均绝对误差。然而，去噪预训练方法仍然有效，并且优于从头训练的模型。特别是，Frad始终优于Coord。有趣的是，Frad在不准确构象下的训练甚至超过了在准确构象下训练的Coord。这些发现验证了Frad是一种高效的预训练模型，即使使用不准确的构象，也使得在更大规模和使用较低准确度的预训练数据集上进行预训练成为可能。

如前所述，Coord方法面临分子分布建模偏差的挑战，这导致采样范围受限和力目标不准确。在本节中，作者将定量验证Frad如何增强采样覆盖率并提高力的准确性，从而实现更优的下游性能。

为评估去噪预训练中力目标的准确性，作者使用皮尔逊相关系数ρ量化估计力与真实值之间的精度。真实力通过一种名为sGDML的监督力学习方法建立。为公平比较，作者将采样和力计算解耦，通过扰动阿司匹林分子的平衡构象来生成样本，采用不同的噪声设置（τ = 0.04, σ = 1；τ = 0.04, σ = 20），从近到远离平衡。结果显示，在所有采样设置中，RN与坐标噪声的混合使用在力准确性上始终优于仅使用坐标噪声。特别是σ = 20和τ = 0.04的配置与真实力场的对齐最佳。由于Cerror随着σ增大而增加，作者选择σ = 2和τ = 0.04作为Frad(RN)的噪声规模。此外，当样本远离平衡时，混合噪声与坐标噪声之间的准确性差距更为明显。

为比较不同噪声类型的采样覆盖率，作者通过施加噪声后所有原子的平均绝对坐标变化定义扰动尺度，并测量其对应的MAE下游性能。结果显示，首先，Coord的低采样覆盖率问题明显，具体表现为其下游性能对CGN的方差敏感，其中τ = 0.04表现最佳，过大或过小的噪声尺度显著降低性能。这一现象可归因于较大噪声尺度导致更多不合理的噪声样本，而较小尺度则产生琐碎的去噪任务，这与其他研究的发现一致。其次，RN缓解了低采样覆盖率的问题，特别是RN的扰动尺度可以显著增加而不失去能力，即使在σ = 20的情况下，Frad(RN)在所有设置中仍优于Coord。最后，更准确的力近似有助于下游任务性能，因为Frad在与Coord共享相似扰动尺度的情况下仍获得了进一步的改进。

本文提出了一种分子预训练框架Frad，以学习有效的分子表示。Frad结合了混合噪声策略，确保力学习的解释能力，并灵活设计噪声，融入化学先验以优化分子分布建模，使得Frad能够在远离平衡状态的低能构象中采样并学习更准确的力。实验结果表明，Frad在力预测、量子化学性质预测和结合亲和力预测任务中均超越了现有的基线模型。此外，Frad对不准确的3D数据具有鲁棒性，显示出更高的力精度和更大的采样覆盖率。未来可通过增加预训练数据量、与其他预训练方法结合，以及针对特定分子类别设计CAN，进一步推动Frad在药物和材料发现领域的应用。

编译|于洲

审稿|王梓旭

Ni Y, Feng S, Hong X, et al. Pre-training with fractional denoising to enhance molecular property prediction[J]. Nature Machine Intelligence, 2024: 1-10.