SchNet|一种适用于分子和材料的深度学习架构

深度学习促进人工智能快速发展，并对语音识别以及生物信息学，物理化学的影响越来越大。一般而言，机器学习特别是深度学习非常适合于表示量子力学相互作用，能够建模非线性势能面或加强对化合物空间的探索。在这里，文章介绍了深度学习架构SchNet，该架构专门设计用于通过使用连续卷积层对原子系统进行建模。文章利用SchNet来预测小分子的分子动力学模拟的势能面和守恒力场，并对c₂₀-富勒烯的量子化学性质进行了示范性的研究，这在常规的从头算分子动力学中是不可行的。

一、介绍

加速发现具有所需性质的分子和材料是计算化学和材料科学长期面临的挑战。然而，在探索广阔的化学空间时，精确的量子化学计算的成本却令人望而却步。近年来，使用机器学习来克服这一瓶颈的研究有所增加，目前，深度学习已经应用于使用图神经网络从分子结构预测性质。然而，由于输入中缺少原子位置，这些仅限于对平衡结构的预测。直到最近，才开发出直接从原子类型和位置学习的预测方法。虽然神经网络通常被认为是一个“黑盒”，但最近人们越来越努力地解释它们的预测，以了解它们是如何运作的，甚至是如何提取科学见解。这可以通过分析经过训练的模型或直接设计可解释的模型来实现。对于量子化学，一些人提出了这样一种具有深度张量神经网络（DTNN）的可解释体系结构，它不仅可以学习原子环境的表示，而且可以从空间和化学上解析量子力学观测。

在此，文章以这项工作为基础，提出了深度学习体系结构SchNet，它允许对复杂的原子相互作用进行建模，以便预测势能面或加速化学空间的探索。SchNet作为DTNNs的一个变体，能够在学习中遵循原子系统旋转和平移不变性以及原子排列不变性。这使得能够准确预测整个构型化学空间，其中原子之间的相互作用是使用连续卷积层建模的。文章证明，这些可以有效地结合周期性边界条件，从而能够准确预测各种大块晶体的形成能。

二、方法

SchNet是早期提出的深度张量神经网络（DTNN）的一个变体，因此共享许多基本构建块。其中包括原子嵌入块、相互作用块和原子能量贡献块。在每一层，原子系统都是通过与周围原子的成对相互作用来细化原子的。在DTNN框架中，相互作用块由张量层来进行建模。SchNet利用连续卷积和卷积核生成网络对相互作用块项进行建模。在下文中，文章将介绍这些组件，并描述它们是构成SchNet体系结构的。有关SchNet体系结构的概述，请参见图1。

图1.chNet架构图（左）和交互块（右）

2.1原子嵌入

原子系统可以用一组具有核电荷Z=(Z₁，...，Z_n )和位置R=(r₁，...，r_n)的n个原子点来唯一地描述。通过SchNet的原子嵌入层，原子由一组特征X^l=(x^l₁，…，x^l_n)来描述，x^l_i∈R^F，特征映射的数量用F来表示，原子的数量用n来表示，当前的层数用l来表示。原子i的表示使用依赖于原子类型Z_i的嵌入进行初始化：

这些嵌入a_z在训练期间随机初始化。它们代表的是一个系统的原子，暂时不考虑其环境的任何信息。

2.2原子层

原子层是密集的层，分别应用于每个原子i的表示x^l_i：

由于权重W^l和偏差b^l在原子之间是共享的，文章的体系结构对于原子的数量仍然是可扩展的。当原子表示通过网络时，这些层会转换它们，并处理通过相互作用层合并的原子环境的信息。

2.3交互块

这里文章用连续滤波卷积层(cfconv)建模，这是离散卷积层的推广，例如用于图像或音频数据。这种泛化是必要的，因为原子不像图像像素那样位于规则的网格上，而是可以位于任意的位置。因此，在传统卷积层中使用的卷积核并不适用。相反，文章需要用一个生成卷积核的神经网络来建模连续卷积。给定R位置R的原子表示X^l，文章得到了原子i与周围所有原子的卷积：

cfconv层和三个原子层构成了相互作用块的残差，文章使用ssp(x)=ln(0.5e^x+0.5)作为整个网络的激活函数。该激活函数保证了ssp(0)=0，并在保证了连续性的同时提高了网络的收敛性。这使文章能够获得光滑的势能面、力场和二阶导数，这是用力进行训练和计算所必需的。

2.4卷积核生成网络

卷积核生成网络决定了原子之间的相互作用如何建模，并可以用来约束模型并包含化学信息。文章选择一个全连接的神经网络，它将从原子i到邻居j的向量作为输入，获得卷积核的值W(r_j−r_i )(见图2(左))。这使得文章能够将已知的分子和材料的不变性纳入模型中。

图2.卷积核生成网络(左)与pbc示例(右)

2.4.1旋转不变性

通过计算成对距离而不是使用相对位置，可以直接包括旋转不变性。文章在高斯分布的基础上进一步扩展了距离

2.4.2周期边界条件

对于具有周期边界条件的原子系统(pbc)，每个原子级特征向量x_i在所有周期重复中都必须是等价的，即对于重复单元元胞a和b有：x_i=x_ia=x_ib。由于卷积的线性，文章能够将pbc直接应用于卷积核，以准确描述原子的相互作用。给定一个用||r_jb−r_ia||_cut覆盖所有原子的卷积核Wl(r_jb−r_ia)，文章得到了卷积。

这个新的卷积核W现在依赖于系统的pbc。文章发现，当卷积核响应的原子数x^l+1_i在截止范围内归一化时，训练更稳定。图2（右）中显示了一些可生成的没有pbc的卷积核，有立方金刚石晶体和六角形石墨。金刚石和石墨的卷积核是单原子卷积核根据各自晶格的叠加，它们反映了晶格的结构。

2.5属性预测

最后，根据所得到的原子表示法预测了一个分子或材料的给定性质P。文章从全连接的预测网络中计算原子贡献P_i（见图1中的蓝色层）。文章分别通过对原子贡献求加或平均来计算最终的预测P。

其中，计算原子间力的公式如下：

2.6训练

文章通过最小化平方损失来训练每个属性目标P的SchNet。

为了训练分子动力学轨迹的能量和力，文章使用了一个联合损失。

其中ρ是能量和力损失之间的权衡。

三、结果

3.1学习分子特性

文章训练SchNet模型来预测QM9数据集的各种性质，该数据集由131k个有机小分子组成，其中最多有9个重原子。文章使用了一个包含10,000个分子的验证集。文章使用T=6的交互块和具有F=64的特征维度，并执行多达1000万个梯度下降参数更新。由于QM9的分子很小，所以文章不使用截断半径。预测误差列于表1中，其中，文章与使用enn-s2s29进行比较。SchNet对极化率α和电子空间范围hR2i的预测在准确性方面明显不足。但SchNet在12个特性中超越了enn-s2s的预测。

表1.QM9数据集上的能量预测的平均误差

图3中显示了T∈{1、2、3、6}交互块的学习曲线，与性能最好的DTNN模型相比。使用T=3交互块的性能最好的DTNN只能优于使用T=1的SchNet模型。文章观察到，使用110k个样本，超过两个交互块的误差仅从T=2交互块的0.015eV略微提高到T∈{3,6}的0.014eV。当用较少的样本进行训练时，差异变得更加显著，而T=6虽然参数最多，但误差最小。此外，该模型需要的收敛时间要少得多，例如，使用110k个样本将所需的次数从T=2的2400次减少到T=6的750次以下的epoch。

图3.QM9数据集下不同参数选择的平均能量预测误差

3.2学习材料的形成能

文章使用SchNet来预测来自Materials Project(MP)的69,640个结构晶体的形成能。它由各种各样的晶体组成。平均绝对误差列于表二。同样，文章使用T=6交互块和具有F=64特征维度的原子表示。文章设置了截断半径rcut=5˚A，将数据随机分为60,000个训练示例，4,500个验证示例，其余数据作为测试集。尽管MP存储库比QM9分子基准更多样化，SchNet能够预测形成能的平均绝对误差为0.035eV/atom。在一个3000个较小的子集训练例子上，SchNet仍然比Faber等人提出的描述符实现MAE达到0.127eV/原子的显著改进。

表2.MP数据集上原子形成能量预测的平均误差

3.3在c₂₀-富勒烯分子动力学中的应用

在MD17基准集上演示了SchNet的准确性之后，文章对C₂₀富勒烯的ML驱动分子动力学(MD) 模拟进行了研究。这个中等大小的分子有一个复杂的PES，在这里，文章使用SchNet对引入核量子效应（NQE）时C₂₀的PES的一些基本性质进行了分析。参考数据是通过在理论的广义梯度近似（GGA）水平上使用DFT运行经典MD生成的，使用Perdew-Burke-Ernzerhof（PBE）交换相关函数和Tkatchenko-Scheffler（TS）方法来解释范德华相互作用。

为了获得准确的能量和力预测，文章首先对给定的参考数据进行广泛的模型选择。文章使用20k 个C₂₀样本作为训练集，4.5k个样本用于提前停止，并报告剩余数据的测试错误。表3列出了交互块数T、原子表示的特征维数F和ρ的各种设置的结果。首先，文章选择模型的最佳超参数T，F，给出ρ=0.01。文章发现T=6和F=128的构型对能量和力都最有效。

对于能量预测，文章在使用ρ=0.1时获得了最好的结果。在这里，文章选择仅力模型作为力场来驱动文章的MD模拟。总的来说，使用SchNet，文章可以执行1.25ns的PIMD，与DFT相比，运行时间减少了3-4个数量级，计算资源要少得多。

表3.c₂₀-富勒烯的能量和力预测的平均误差

4、结论

文章提出了SchNet作为一种通用的量子化学深度学习体系结构，从不同分子和材料数据集的特性预测势能面和力场。SchNet能够直接建模周期边界条件。除此之外，文章已经在PBE+vdWTS理论水平上对富勒烯C₂₀进行了一个示例性的分子动力学研究，这在常见的DFT方法中是不可行的。这些结果可以指导未来的工作，如研究更大的分子和周期系统，以及进一步拓展可解释的深度学习架构，以协助化学研究。