ForceNet：一个用于大规模量子计算的图神经网络

科学家们手上已经有了大量的原子模拟数据，如果以此来开发快速而精确的机器学习模型，从而模拟昂贵的基于物理的计算。评判这类模型好坏的关键量是原子力，在各种模型中，图神经网络（G 满足物理约束，现有的模型必须在计算效率和模型表达能力之间做出权衡。

一篇新的论文探讨一种替代方法：通过不施加明确的物理约束，我们可以在保持计算效率的同时，灵活地设计表达模型。通过使用增强后的基于物理的数据来训练模型，物理约束将被隐含在模型中。

为了评估这种方法，论文重新设计了一个可表达的 GNN 模型：ForceNet，并且它支持后续的功能拓展。它能够比目前最先进的基于物理的 GNN 更准确地预测原子力，同时训练和预测所需要的时间也更短。将这个模型应用于 OC20（一庞大的量子物理计算数据集），最终的实验结果表明，这项工作将会为模拟量子计算开拓新的途径。

复杂物理系统建模的新方法

科学计算产生了大量基于物理的数据，这为机器学习 (Machine Learning) 精准地建模复杂的物理系统，提供了巨大的机会。在大数据上训练得到的机器学习模型可以精确地预测数量级，速度还比原始的基于物理的计算要快。

估算量子力学系统中的原子力是具有极高的实际意义的。因为每个系统都需要几个小时的时间去运行计算，潜在的需求也很多，如原子力可以用于不同的化学应用，如结构弛豫、分子动力学、结构分析以及过渡态计算。

目前，预测原子力最先进的方法，是 Gilmer 在 2017 年提出的基于物理信息传递的图神经网络 (GNNs)，最具代表性的模型就是 SchNe 和 DimeNet。这些 GNN 首先以旋转不变量的方式预测整个系统的能量，然后通过对原子所在位置的能量求导来预测每个原子的原子力。这种结构设计下的 GNN 计算出的原子力，遵循旋转协方差和能量守恒的基本物理规则。

然而，在同时满足这些物理规则的条件下设计有效的 GNN 是并不简单的。例如，SchNet 计算的速度很快，这是因为该模型在信息传递中只使用原子距离来确保其能量预测的旋转不变性，但 SchNet 却未能明确的反应 3D 结构，这导致 SchNet 的性能出现次优泛化。

在 Klicpera 提出的 DimeNet 和 DimeNet++ 中，网络在信息传递时额外地捕获键角信息，弥补了 SchNet 中的缺失，但为了确保旋转不变性，在涉及原子三元组的计算时，网络消耗了大量的计算成本。因此，DimeNet 需要大量的计算才能扩展到大型数据集 (图 1 (左))：DimeNet 总共用了 1600 个 GPUdays 才完成训练。而且，即使是 DimeNet 也无法模拟出扭转角（torsion angles）这一重要的物理特征，无法在其信息传递中捕捉到完整的 3D 结构。

本项研究，以 Bapst 的基于图网络的模拟器 (GNS) 为框架提出了一种替代方法。具体地说，就是通过不在模型架构中显式地施加物理约束，可以灵活地设计更富有表现力的 GNN 模型。

本文设计的 GNN 模型，以可伸缩的灵活的方式精确捕捉 3D 原子结构，最终将模型的名字命名为 ForceNet。研究团队使用了精心选择的基础函数和非线性激活函数，使其成为一种有效的消息传递架构。

训练 ForceNet，使用了 OC20 数据集 (Chanussot，2020 年)：一个最近引入的大规模量子物理计算数据集，包含 2 亿多万个大原子结构 (20 - 200 个原子)，从中能应用于能源应用的新催化剂 (图 2)。数据集以 Density Functional Theory (DFT) 的量子计算为基础，构造了前所未有的 7000 万 CPU 小时的计算量，这比传统的量子物理的数据集 QM9 和 Alchemy 的 20 倍还要多。正是这样夸张的计算量使其成为深度学习的理想数据集。

除此以外，与许多现有的量子数据集不同，OC20 提供了分子的非平衡结构，即具有非零原子力的三维结构，这正是测试 ForceNet 模型的不二选择。

尽管 ForceNet 没有任何显式的物理约束，但使用如此大的数据集进行训练，ForceNet 同样能够达到很高的精度 (图 1 (左))，甚至超过了基于物理的 GNN。此外，ForceNet (ForceNet-large) 最终的预测误差与顶尖的 DimeNet++ (DimeNet++-large) 相当，同时计算所消耗的时间只有后者的六分之一 (图 1 (右))。总的来说，与 DimeNet++ 相比，ForceNet-large 实现了更准确的原子力预测，同时训练和预测的时间更短 (图 1 (左) 和 (右))。

为了解释 ForceNet 的设计选择，团队对 ForceNet 的每个架构组件都进行了大量的研究。结果发现，三维原子相互作用的边缘级计算是模型中表现最突出的，也是对模型性能贡献最大的模块。总之，即使没有显式的物理约束，可扩展的表达型 GNN 在复杂物理系统建模方面也表现出很好的性能，ForceNet 就是证明，这为未来的研究开辟了新的方向。

图 1：不同模型 S2F (原子力预测) 的性能，同时考虑计算效率。

图 2：OC20 数据集采样系统的说明。每个体系都由吸附质 (表面的小分子) 和催化物质 (吸附质下面的大网格状分子) 组成，并在水平方向上无限重复。ForceNet 旨在有效地预测每个原子的原子力。

ForceNet 的具体架构

为了提高 ForceNet 的旋转协方差预测准确率，本节将介绍对有效的数据增强策略。ForceNet 的输入是一个原子结构，即一组原子及其三维空间结果 (图 2)。输出是每个节点的三维向量，表示预测的 (x, y, z) 三个方向的原子力。

图 3：Eqns 中 ForceNet 使用的消息 mst (从 atom s 到 atom t) 的模型图。

ForceNet 将原子表示为 GNN 中的节点，原子间的相互作用（原子力）表示为边。节点输入分别是指定原子的原子序数和其他属性 (9 维向量) 的特征。GNN 的边缘是由相邻原子的半径图构造的。设  为以目标原子 t 为中心，距离小于长度 c 的一组相邻原子，一个原子平均有 35 个相邻原子。s∈  为一条从源原子 s 到目标原子 t 的有向边。设  ∈R3 为它们的相对位移，即从原子 s 指向原子 t 的矢量。

ForceNet 遵循 GNS 框架的编码器 - 解码器架构。编码器使用可伸缩的迭代消息传递来计算节点嵌入  ，以此来捕获每个原子周围的 3D 结构，解码器使用多层感知器 (MLP) 直接从这些嵌入  中预测每个原子的力。编码器将  更新为:

  （1）

其中，消息  和  求和，并通过函数  传递，这是一个 1 层的批归一化 MLP 网络。公式 (1) 遵循标准的 GNN 嵌入更新公式，并添加一个残差连接  。我们在 3.1.1 节中定义了成对的消息  和  。

解码器的计算采用最后一层 K 的节点嵌入  ，  其中  为原子 t 的三维力，  为批处理归一化的 1 隐层 MLP。

ForceNet 的关键是它的编码器，特别是可扩展的消息计算，它可以有效地捕捉非线性的复杂的 3D 原子相互作用，从而预测原子力。

接下来，将介绍消息计算架构中的几个关键组件。

有条件的卷积滤波器

模型首次提出了有条件的卷积滤波器这个概念，它是对 SchNet 中可扩展连续卷积滤波器的一个简单而有效的扩展。最初的连续卷积滤波器利用相邻原子之间的距离信息来计算滤波，然后将结果应用于源原子的嵌入。然而，这种方法有一系列的局限性，特别是从以能量为中心的模型迁移到以力量为中心的模型时，表现的尤为明显。

首先，为了保证旋转不变的能量预测，SchNet 连续滤波仅使用原子距离作为输入边缘特征。这导致角动量的信息会在消息传递时丢失。更严重的是，用这种方法产生的嵌入节点是旋转不变的，但力是需要与分子系统一起旋转。此外，SchNet 的滤波器不依赖于源原子和目标原子：由于不同原子拥有不同的电子性质，即使它们被放置在相似的距离，它们的原子类型的不同也会导致原子之间力的显著差异。因此，SChNet 中的滤波器可能不足以对复杂的非线性原子间的相互作用进行建模。

使用 E 维边缘特征  (如下所述) 来解决这些问题，该特征编码了有指向的旋转协变信息，并同时约束源节点  和目标节点  信息:

（2）

与现有 GNS 模型的比较。   我们注意到滤波器  本身已经包含了源节点信息  ，并且可以像现成的 GNS 模型那样直接聚合。但是我们发现，显示地将  应用在公式 (2) 中  的元素点积之上可以显著提高性能。ForceNet 和现有的 GNS 模型之间还有其他微妙但重要的区别，例如基函数的使用和非线性激活的选择。经检验，与现有的 GNS 模型相比，ForceNet 中这些谨慎的架构设计选择显著地提升了性能表现。

执行旋转协方差。需要注意的是，在  中使用  会导致模型不总是旋转不变的，也不一定是协变的。在第 3.2 节中，我们建议通过旋转数据增强来训练模型，以此实现物理约束的功能。

MLPS 中的非线性表达

这一节将会介绍架构的最后一个部分，虽然这一部分很简单，但却至关重要。原子的相互作用十分复杂，因此非线性激活函数的选择在建立非线性模型中起着关键作用。ReLU 激活 (gloria et al.， 2011) 被广泛应用于许多深度学习的模型中，包括现有的 GNS 模型。然而，ReLU 在建立原子力模型的方面却并不理想，因为它将模型的输出塑造为具有尖锐边界的分段线性超平面。但在理想情况下，需要的非线性激活函数是将输出变得平滑且富有表现力的。

在多种激活函数的进行选择，如 Tanh、Leaky-ReLU、SoftPlus 、shiftsoftplus 和 Swish 中，发现 Swish。如图 4 所示，Swish 使输出更平滑，并且对负输入具有非零激活的功能。用 Swish 替换 ReLU 可以显著地提高预测精度，同时让所有基函数都保持有可伸缩性。

旋转数据增强

本文使用随机旋转数据增强来提升模型预测的旋转协方差。具体来说，将整个系统和每个原子的力分别随机旋转相同的角度，并在这样的情况下让 ForceNet 根据旋转系统的信息来预测旋转后的原子力。

当目标物理系统有典范轴时，旋转增强尤其有效。在现实应用中，有许多这样的系统，例如，大多数物质型分子系统都有垂直于物质表面的典范轴 (如图 2)。在宏观层面上，大多数物理系统都有指向地球引力方向的自然典范轴。在有这样一个典型轴的情况下，旋转增强只需要沿着单个轴进行就能有效地提高学习的效率。

在这项工作中，使用具有垂直于材料表面典范轴的原子结构来训练 ForceNet，明确了这个属性，只需要在垂直轴方向进行旋转增强。在这种策略下，ForceNet 就可以使用大量的数据进行训练，从而学习到接近旋转协方差的效果。

实验结果与总结

最终，将模型应用到 OC20 数据集上，评估 ForceNet 在预测原子力方面的性能。OC20 是一个量子物理计算的海量数据集，主要是探索催化剂的非平衡原子结构。

以下为论文中实验结果：

总而言之，这项工作通过精心设计信息传递架构、在拥有海量数据的数据集上训练模型，并使用基于物理的数据增强来完善了 ForceNet。通过实验证明，在没有直接添加物理约束的情况下，侧重于力中心的 GNN 模型在预测原子力的方面，比最先进的能量中心 GNN 模型拥有更高的准确性，同时在训练和推理方面所用的时间也更短。

该工作有望为未来的研究开辟了许多方向：(1) 将同样的原理应用于 OC20 的其他预测任务 (例如，预测每个系统的能量) 或其他应用领域。(2) 将物理知识融入 ForceNet，从而提高其泛化性能。(3) 在保持 ForceNet 性能的前提下，提高它的计算效率，类似于 DimeNet ++ 在 DimeNet 的基础之上做出的改进。(4) 增加 ForceNet 的计算资源。

论文：

ttps://arxiv.org/abs/2103.01436v1

 往期推荐 

丨语言模型能知道一头大象有多重吗？

丨麻省理工学院学者万字长文：计算机作为一种通用技术的衰落

丨华为 AI 水军被推特封禁的背后，谁应为 AI 应用的边界负责？

丨果蝇能学会 Word Embedding 吗？

关于数据实战派

数据实战派希望用真实数据和行业实战案例，帮助读者提升业务能力，共建有趣的大数据社区。