这周读了一篇有关分子性质的文章《Geometry-enhanced molecular representation learning for property prediction》，文章于2022.2.7发表在 Nature Machine Intelligence 期刊上，期刊属于计算机 人工智能1区，3年平均IF为15.508。文章的创新点在于将3D信息（键长、键角、原子之间的距离）应用到分子表示当中。

1 Introduction

1.1 背景

随着DNN的发展，分子表征学习展现出了巨大的优势，其在分子性质预测方面的应用也越来越广泛，它将原子和键的拓扑结构视为一个图形，并将每个元素的信息传播到其邻域。
最近的研究开始在自监督方法中使用大规模未标记分子来预先训练分子表征，然后使用少量标记分子来微调模型。但现有的自监督技术只考虑分子的拓扑信息，忽略了分子的几何结构（3D结构）。

1.2 本文工作

本文提出了一种新的几何增强分子表征学习方法（GEM）。
①首先，为了使信息传递对几何体敏感，本文通过设计基于几何体的GNN体系结构（GeoGNN），同时对原子、键和键角的影响进行建模。
该体系结构由两个图组成：第一个图将原子视为节点，键视为边，而第二个图将键视为节点，键角视为边。
②其次，对GeoGNN进行预训练，从具有粗糙三维空间结构的大分子中学习化学定律和几何，设计各种几何级别的自监督学习任务。
③最后，为了验证所提出的GEM的有效性，在15个分子性质预测基准上将其与几个最先进的（SOTA）基线进行了比较，其中GEM获得了14个SOTA结果。

2 GEM框架

2.1 Geometry-based GNN architecture (GeoGNN)（对应3.2）

本文提出了GeoGNN，它通过模拟 原子-键-角度 关系来编码分子的几何结构信息。如图，左侧为atom–bond graph G=(V,ε)，右侧为bond–angle graph H=(ε,A)，其中V是原子集，ε是化学键集，A是键角集。

将 xu 作为原子u的初始特征，xuv 作为边(u,v)的初始特征，xuvw 作为键角(u,v,w)的初始特征。
将原子键图G和键角图H，原子特征、键特征和键角特征作为GeoGNN的输入，GeoGNN迭代学习原子和键的表示向量。对于第k次迭代，原子u和键(u,v)的表示向量分别用 hu 和 huv 表示。
为了连接原子键图G和键角图H，键的表示向量被视为G和H之间的通信链路。
step1：在键角图H中聚合相邻 键与键角 的信息来学习键的表示向量。
step2：在原子键图G中聚合相邻 原子与键 的消息来学习原子的表示向量。
step3：最后，通过聚合原子的表示来获得分子表示 hG。
（公式详见方法部分）

2.2 Geometry-level self-supervised learning tasks

为了进一步提高GeoGNN的泛化能力，本文提出了三个几何级的自监督学习任务来预训练GeoGNN，分别是：
（1）键长预测；
（2）键角预测；
（3）原子距离矩阵预测。
其中，键长和键角用于描述局部空间结构，而原子距离矩阵用于描述全局空间结构。

2.2.1 局部空间结构

键长和键角是最重要的分子几何参数。其中，键长是分子中两个连接的原子之间的距离，反映原子之间的键强度；键角是两个连续键（包括三个原子）的角度，描述分子的局部空间结构。为了学习局部空间结构，本文构建了预测键长和键角的自监督学习任务。
预测局部空间结构的任务可以看作是一个节点级的自监督学习任务。
step1：对于一个分子，随机选择15%的原子，对于每个选定的原子，提取该原子的一跳邻域，包括相邻的原子和键，以及该选定原子形成的键角；
step2：在一跳邻域中屏蔽这些原子、键和键角的特征。在GeoGNN的最终迭代中，提取的原子和键的表示向量 用于预测提取的键长与键角。
本文设计了一个回归损失函数来惩罚预测的键长/键角和标签之间的误差。（公式详见方法部分）

2.2.2 全局空间结构

除了学习局部空间结构的任务外，本文还设计了学习全局分子几何的原子距离矩阵预测任务。
本文根据原子的三维坐标为每个分子构造原子距离矩阵，然后预测距离矩阵中的元素。注意，对于具有相同拓扑结构的两个分子，相应原子之间的空间距离可能会有很大差异。（损失函数在方法部分）

因此，对于一个分子，本文没有将预测原子距离矩阵作为一个回归问题，而是将其作为一个多分类问题，通过以相等的步幅将原子距离投影到30个二进制表示中，也就是one-hot。有关设计损失函数的详细信息在方法部分。
为了预训练GeoGNN，本文通过总结相应的损失函数来考虑每个分子的局部空间结构和全局空间结构。
预测键长与原子距离的区别：预测键长的任务可以看作是预测原子距离任务的一个特例，但键长预测更关注局部空间结构，而原子距离更关注全局空间结构的分布。

3 方法

GNN是一种消息传递神经网络，给定节点v，它在第k次迭代中的表示向量 hv(k) 为：

最后使用READOUT函数来整合节点信息，得到图的表示 hG ：

3.2 GeoGNN(对应2.1)

step1：在键角图H中聚合相邻 键与键角 的信息来学习键的表示向量。
键，也就是边(u,v)，在第k次迭代中的表示向量 huv(k) 为：

step2：在原子键图G中聚合相邻 原子与键 的消息来学习原子的表示向量。
原子u在第k次迭代中的表示向量 hu(k) 为：

step3：在最后一次迭代的时候，通过聚集原子的表示来获得分子表示 hG，hG 用来预测分子的性质。

3.3 Geometry-level self-supervised learning tasks

3.3.1 局部空间结构

局部空间信息的自监督任务用于学习两个重要的分子几何参数：键长和键角。本文设计了一个回归损失函数来惩罚预测的 键长/键角 与 标签之间的误差，其定义如下：

3.3.2 全局空间结构

全局空间信息的自监督任务用于学习 所有原子对 之间的原子距离矩阵，距离矩阵中的每个元素都是两个原子之间的三维距离，用 duv 来表示分子中原子u和原子v之间的距离，损失函数如下：

4 实验

4.1 Basic Setting

①数据集：
预训练：Zinc15
分子性质预测：MoleculeNet

②GNN结构：使用GIN中定义的聚合函数与组合函数，然后加上残差连接、层归一化、图归一化 来进一步提高性能，使用平均池化作为ReadOut来获取图表示。

③评估指标：
分类：ROC-AUC
回归：RMSE 与 MAE

4.2 实验结果——①所有任务的总体性能

表格：
①最好的结果用黑体表示，括号内是偏差
②ToxCast 与 PCBA 上的N-Gram太耗时，无法按时完成
③回归的结果越小越好（因为是差值），分类的结果越大越好

结果解读：
①在15个数据集的14个中都取得了最先进的结果。
②在回归任务上，与之前最先进的结果相比，平均提升为8.8%；分类任务上的提升为4.7%。
③在分类任务上提升更大，可能是因为回归数据集侧重于预测 与分子几何形状高度相关的 量子化学性质。

4.2 实验结果——②回归任务中不同GNN结构的性能

结果：①GeoGNN 在所有回归数据集上都大大优于其他 GNN 架构，与以往方法的最佳结果相比，整体相对提高了 7.9%；
②因为即使模拟了原子的三维坐标，GeoGNN 也包含几何参数。

4.2 实验结果——③回归任务中不同预训练策略下GeoGNN的表现

为了研究所提出的几何级自监督学习任务的效果，采用不同类型的自监督学习任务对回归数据集上的GeoGNN进行预训练。

表格：“无预训练”表示无预训练的GeoGNN网络，“Geometry几何”表示我们提出的几何级别任务，“Graph图形”表示预测分子指纹的图形级别任务，“Context上下文”表示预测原子上下文的节点级别任务。

结果：①有几何级别任务的方法比没有几何级别任务的方法要好。
②“几何”在回归任务中的表现优于“几何+图形”，这可能是因为分子指纹与回归任务之间的联系比较弱。

5 结论

总结：现有的分子性质预测的预训练方法没有利用由键、键角和其他几何参数描述的分子几何结构，因此本文提出了基于几何的GNN结构（GeoGNN），并采用多几何层次的自监督学习方法来获取分子的空间知识。在多个数据集上进行了实验，GEM的表现优于其他基准方法。

创新点：提出并应用了基于键长与键角的自监督学习。