图神经网络与分子表征：7. LEFTNet

在执行性质预测任务时，我们需要考虑两个问题：1. 如何正确的将图结构进行编码？2. 如何汇聚编码信息预测整个分子的任务？
LEFTNet 就是通过回答上述问题来进行模型设计的。
原文地址

算法设计

原文中，作者定义了三个图同构问题：

1. 树同构：任意邻接的两原子间的距离一致即可
2. 三角同构：任意邻接的三原子构成的三角形一致
3. 子图同构：任意两个三角形的相对位置保持一致

作者认为，一个好的编码器能够甄别上述所有同构问题，就像一台高分辨率的相机一样。为了方便理解，作者给出了两个案例：

在图（a）中，作者展示了一种，满足树同构，但不满足三角同构的情况。
在图（b）中，作者展示了一种，满足三角同构，但不满足子图同构的情况。
有了模型设计方向，作者开始讲述设计思路：
首先，为了满足树同构的要求。作者认为，只要正确抽取原子间距离即可，所有的 3DGNN 均满足。
其次，为了满足三角同构的要求。作者认为，三角同构本质上是局域完备性问题。作者设计了一个局域坐标系，通过将原坐标映射到局域坐标，即可满足局域完备性的要求，同时还能满足等变要求。此处直接采用了作者上一篇文章 ClofNet。
注：能够满足局域完备性的模型很多，可以参考我之前的博客
最后，为了满足子图同构的要求，作者认为，不同局域坐标系相对位置不一样，这导致局域信息汇聚过程中遭受信息损失。在信息汇聚过程中有效考虑不同坐标系的相对位置，能够有效解决这一问题。

如上图所示，点 b 和点 c 分别建立了局域坐标系，在这两个局域坐标系下，cluster（团簇）b 和 cluster c 进行了有效映射。b 处的局域信息如何才能有效和 c 处的局域信息融合呢？
作者指出，通过二者共同的邻居 a，以及两个局域坐标系的相对位置，能够实现信息的转换融合。因此，模型设计时需要考虑两点：

1. 编码时需要纳入尽可能完备的局域信息；
2. 信息汇聚时需要考虑不同局域坐标系的相对位置；

算法实现

一共两处 github 仓库与 LEFTNet 的实现有关：m2hub, leftnet

二者差别不大，此处以 m2hub 中的为例。

作者在原文中给出了模型框架图（上图）和算法流程（下图）：

实际代码实现与原文存在些许区别。按原文的模型框架图，LEFTNet 是向量、标量双通道模型，与 PAINN 一致。
在消息传递模块中，模型会首先通过 LSE 抽取子图信息，并将其抽象成邻居消息的权重 Aij.
邻居消息就是边的 embedding，通过与权重进行糅合，有效融入局域信息。
最后，在消息的聚合阶段，通过考虑不同局域坐标系的相对位置，实现子图级的消息融合。
下面我们读一下代码：

1. 抽取局域信息的 LSE 模块
   该模块在 interaction 之前，属于数据预处理：
   
2. 消息传递主函数
   
   实际实现过程中，原文提到的 FTE 模块进行了简化，取消了node_frame的介入，因此可以看做是简单的 scalar vector infusion
3. 消息传递模块
   此处是直接调用的 torch geometric 的消息传递类
   原文中消息是由三个成分构成，实际代码实现中仅使用了前两项
   
   消息传递模块如下：
   
4. 标量、向量融合阶段

对应原文中的 FTE（进行了简化）

核心思路是向量标量化，与原标量信息融合。融合标量再向量化，与原向量信息进行融合。
使用类似融合方法的模型很多，例如：
so3krate 和 torchmd-net, BAMBOO 中的 GET layer，如下图所示：

写在最后

LEFTNet 这篇论文行文晦涩难懂，这是和其他模型文章最大的区别。大量的定理证明展现了作者扎实的数学功底。在 LEFTNet 之后，西班牙团队很快发展出了 torchmd-net （有名 tensornet, neruips 2023），torchmd-net 在纳入电荷信息后升级到了 torchmd-net 2。字节团队在 torchmd-net 2 之上搭建了 BAMBOO。这些论文相对容易理解，感兴趣的读者可以拓展阅读。