今天给大家介绍一篇关于分子表征学习的文章。分子表征学习（MRL）旨在将分子嵌入到实向量空间中。然而，现有的基于SMILES（简化分子线性输入系统）或GNN（图神经网络）的MRL方法要么以SMILES字符串作为输入，难以编码分子的结构信息，要么过度强调GNN结构的重要性，而忽视了其泛化能力。因此，作者提出使用化学反应来协助学习分子表征，其核心思想在于保持分子在嵌入空间中的化学反应的等价性，即强制让每个化学方程式的反应物嵌入和生成物嵌入的总和相等，该限制在保持嵌入空间的有序性和提高分子嵌入的泛化能力中被证明是有效的。此外，该模型可以使用任何GNN作为分子编码器，与GNN结构无关。实验结果表明，这种方法在各种下游任务中都达到了最佳性能，超过了最佳基线方法。
1.引言
在化学中，一般用IUPAC命名法、分子式、结构式、骨架式等形式来表示一个分子，然而这些形式最初是为人类而不是计算机设计的。为了便于计算机理解和利用分子，MRL被提出，MRL将分子映射到低维空间，表示为稠密向量。分子的学习向量（又称嵌入）可用于广泛的下游任务，如化学反应预测、分子性质预测、分子结构预测等。
研究人员已经提出了许多MRL方法，其中大部分都将SMILES字符串作为输入，尽管这些语言模型功能强大，但它们却很难从SMILES中学习到分子的原始结构信息，因为SMILES是分子结构的一维线性化。而GNN在处理分子表示中局限于设计新颖精致的GNN结构，忽略了MRL的本质——泛化能力。这激励了作者探索GNN架构之外的方法。
在本文中，作者使用化学反应来协助学习分子表征，提高其泛化能力。化学反应通常由化学反应方程式表示，反应物在左侧，生成物在右侧，例如乙酸和乙醇的费希尔酯化反应。作者的想法是保持分子在嵌入空间的等价性。更重要的是，当分子编码器是以总和作为readout函数的GNN时，该模型能够自动隐式地学习同一类别内一组化学反应的反应模板，学习反应模板是提高分子表征泛化能力的关键。
作者将该模型称为MoLR（chemical-reaction-aware molecule embeddings，化学反应感知的分子嵌入）。并且将可视化分子嵌入，表明它们能够编码反应模板和几个关键的分子属性，如分子大小和最小环的数量。
2.方法
结构分子编码器
分子图用G = (V, E)表示，其中V = {a1, …}是非氢原子的集合，E = {b1, …}是键的集合。在这项工作中，每个原子ai都有初始特征向量，该向量包括四个one-hot向量，每个one-hot向量表示一种原子属性，四个one-hot向量分别表示四种原子属性，分别为元素类型、电荷、原子是否是芳香环以及附着的氢原子数，然后每个one-hot向量有一个额外的未知条目，以在推理过程中处理未知值。此外，每个键bi都有一种键类型，例如单键、双键，根据作者的实验，键的类型不会持续改善模型性能，因此没有将键的类型作为输入。
作者使用GNN作为基础模型，典型的GNN遵循领域聚合策略，通过聚合原子的邻居和自身的表示来迭代更新原子的表示。GNN的第ｋ层是：
这里，是原子ai在第k层的表征向量，N(i)是与ai相连的原子集合，K是GNN的层数，这个聚合函数的选择是设计GNN的关键。
最后，readout函数用于聚合最后一个GNN层输出的所有结点表示，以获得整个分子的表示hG：
readout函数可以是简单的置换不变函数，例如求和和取平均，也可以是更复杂的graph-level pooling算法。
图1：(a) GNN编码器处理脯氨酸分子，省略了氢原子；(b) 丙酸和丙醇的费希尔酯化反应，以及模型学习到的相应反应模板。反应中心用橙色表示，距离反应中心1或2的原子用浅橙色表示；(c)小批量化学反应对比损失的图示，dij是嵌入和之间的欧氏距离。
保持化学反应等价性
如前所述，作者希望在分子嵌入空间中保持反应物和生成物的等价性，即：
其中，R是反应物集合，P是生成物集合。这个简单的限制能有效提高分子嵌入的质量，下面的命题将会证明化学反应中的等价关系是等式(3)约束下的等价关系：
对于一个分子的集合M，M的所有子集2的M次方可以根据等价关系划分为等价类，一个等价类中所包含的所有分子的嵌入总和应该相等。因此，等式3形成了一个线性方程组，化学反应等价性对基础分子的嵌入施加了更加强大的约束，分子嵌入的可行解将更加稳健，整个嵌入空间将更加有序。
然后作者提出反应中心的概念，反应中心被定义为将反应物转换为生成物所需的一组最小图形编辑，鉴于反应中心的概念，有以下命题：
命题2表明，反应物嵌入和生成物嵌入之间的剩余仅取决于距离反应中心小于K跳的原子。如图1 (b)所示，剩余物完全取决于反应中心（橙色）以及距离反应中心1或2的原子（浅橙色）。这意味着，方程
适用于任何官能团R1和R2，这称为反应模板，它概括了同一类别内的一组化学反应，反应模板可以轻松地适应训练集中没见过但符合已知反应类型的反应。
训练模型
在这里作者使用了一个类似于基于小批量的对比学习框架。对于小批量的数据B = {R1→P1,R2→P2, …}，首先使用GNN编码器处理B中所有的反应物和生成物，并得到它们的嵌入物，匹配的反应物-生成物对(Ri, Pi)被称为正对，其嵌入差异将最小化，而不匹配的反应物-生成物对(Ri, Pj)(i≠j)，其嵌入差异将最大化。为了避免总损失被负对主导，作者使用基于边距的损失：
其中γ是一个边距超参数。因此，可以使用基于梯度的优化方法（如随机梯度下降SGD）最小化上述损失来训练模型。
3.实验
化学反应预测
作者使用USPTO-479k数据集，每个反应实例包含5种反应物和1种生成物。化学反应预测的任务可表述为一个排序问题。在推断阶段，给定化学反应的反应物集R，将测试集中的所有生成物视为候选者C，根据反应物嵌入hR和候选生成物嵌入hC之间的L2距离对所有生成物进行排序，和真实值比较以计算指标。结果如表1所示，很明显，MoLR的方法都显著优于基线方法。
表1：USPTO-479k数据集上的化学反应预测结果。最佳结果以粗体突出显示，基线的最佳结果以下划线突出显示。
案例研究 作者选取了前20个反应作为案例研究，表2展示了两个反应的结果，结果表明MoLR模型在预测生成物时非常精确。
表2：在USPTO-479k数据集上的案例研究
现实场景生成物预测中的多选择问题 每个问题给出反应的反应物，从4或5个选项中选出正确的生成物，结果如图2所示，MoLR优于基线方法。
图2：现实场景生成物预测中的多选择问题的结果
分子性质预测
作者在5种数据集中测试MoLR，每个数据集包含数千个SMILES分子。预测的AUC结果如表3所示，MoLR在4个数据集中表现最好，作者将MoLR在分子性质预测方面的优异性能归因于，MoLR在USPTO-479k上预训练，因此根据命题2对反应中心敏感。注意，反应中心通常由化学活性官能团组成，这些官能团对确定分子性质至关重要。
图编辑距离预测
表4：分子性质预测的AUC结果
图编辑距离（GED）是两个图之间相似性的度量，定义为将一个图转换为另一个图的最小编辑操作数。这项任务根据两个分子图的嵌入来预测它们之间的GED，目的是显示学习到的分子嵌入是否能够保持分子之间的结构相似性。此外，计算GED是NP-hard问题，这项任务也能近似作为计算GED的算法，结果如表4所示。
表5：在QM9数据集上GED预测的RMSE结果，最佳结果加粗显示
嵌入可视化
图4：在BBBP数据集上可视化分子嵌入空间
在图4a中，分子根据渗透性的性质着色，我们发现了两个非渗透性分子群落，这表明molr可以捕获感兴趣的分子性质。在图4b中，根据BBBP数据集中随机选择的分子（编号1196）的GED对分子进行着色，显然，与1196号分子结构相似的分子（橙色）在嵌入空间中很接近，而与1196号分子结构不同的分子（红色）在嵌入空间中相距很远，这表明molr能捕捉分子的结构相似性。在图4c中，分子根据其大小着色，嵌入空间被完美地分割为小分子区域（上部）和大分子区域（下部）。
换句话说，二维嵌入空间的纵轴表征分子大小。惊讶的是，我们发现横轴与分子中最小环的数量有关：如图4d所示，没有环的分子（蓝色）仅存在于左簇中，含有一个环的分子（黄色）仅存在于左簇和中间簇中，含有两个环的分子（橙色）基本上位于中间簇，而右边簇主要由含有3个及以上环的分子（红色）组成。
图3：醇氧化和醛氧化反应的可视化
作者还以醇氧化和醛氧化为例说明MoLR编码化学反应，结果如图3所示，蓝色箭头的长度大约是相应的红色或橙色箭头的两倍，这正是因为(CH2OH)2/(CH2CHO)2具有两个待氧化的羟基/醛基。
4.结论和未来工作
在这项工作中，作者使用GNNs作为分子编码器，并使用化学反应来协助学习分子表征，强制让反应物的嵌入总和等于生成物的嵌入总和。该模型能够学习反应模板，这是提高模型泛化能力的关键。该模型能够胜任广泛的下游任务，可视化的结果表明，学习到的嵌入是有组织的和反应感知的。
作者指出了今后工作的四个方向。首先，环境条件也是化学反应的一部分，需要考虑进去；第二，值得研究如何显式输出学习的反应模板；第三，值得研究如何在嵌入空间中区分立体异构体，因为现有模型不能处理立体异构体；最后，加入额外的信息（如分子的文本描述）来协助学习分子表征也是一个很有希望的方向。
