13.图神经网络 - 五、分子指纹GCN - 《AI算法工程师手册》

五、分子指纹GCN
- 5.1 模型
- 5.2 实验

五、分子指纹GCN

在材料设计领域的最新工作已经将神经网络用于材料筛选，其任务是通过学习样本来预测新型分子的特性。预测分子特性通常需要将分子图作为输入，然后构建模型来预测。在分子图中顶点表示原子，边表示化学键。

这个任务的一个难点在于：输入的分子图可以具有任意大小和任意形状，而大多数机器学习模型只能够处理固定大小、固定形状的输入。

目前主流的方法是通过hash 函数对分子图进行预处理从而生成固定大小的指纹向量fingerprint vector，该指纹向量作为分子的特征灌入后续的模型中。

论文《Convolutional Networks on Graphs for Learning Molecular Fingerprints》 提出了分子指纹GCN 模型，该模型用一个可微的神经网络代替了分子指纹部分。

神经网络以原始的分子图作为输入，采用卷积层来抽取特征，然后通过全局池化来结合所有原子的特征。这种方式使得我们可以端到端的进行分子预测。

相比较传统的指纹向量的方式，我们的方法具有以下优势：
- 预测能力强：通过实验比较可以发现，我们的模型比传统的指纹向量能够提供更好的预测能力。
- 模型简洁：为了对所有可能的子结构进行编码，传统的指纹向量必须维度非常高。而我们的模型只需要对相关特征进行编码，模型的维度相对而言低得多，这降低了下游的计算量。
- 可解释性：传统的指纹向量对每个片段fragment 进行不同的编码，片段之间没有相似的概念。即：相似的片段不一定有相似的编码；相似的编码也不一定代表了相似的片段。
  
  我们的模型中，每个特征都可以由相似但是不同的分子片段激活，因此相似的片段具有相似的特征，相似的特征也代表了相似的片段。这使得特征的representation 更具有意义。

5.1 模型

5.1.1 圆形指纹算法

分子指纹的最新技术是扩展连接性圆形指纹 extended-connectivity circular fingerprints:ECFP 。ECFP 是对Morgan 算法的改进，旨在以无关于原子标记顺序atom-relabling的方式来识别分子中存在哪些亚结构。

ECFP 通过对前一层邻域的特征进行拼接，然后采用一个固定的哈希函数来抽取当前层的特征。哈希函数的结果视为整数索引，然后对顶点 feature vector 在索引对应位置处填写 1 。
- 不考虑hash 冲突，则指纹向量的每个索引都代表一个特定的亚结构。索引表示的亚结构的范围取决于网络深度，因此网络的层数也被称为指纹的“半径”。
- ECFP 类似于卷积网络，因为它们都在局部采用了相同的操作，并且在全局池化中聚合信息。
ECFP 的计算框架如下图所示：
- 首先通过分子结构构建分子图，其中顶点表示原子、边表示化学键。
- 在每一层，信息在邻域之间流动
- 图的每个顶点在一个固定的指纹向量中占据一个bit。
  
  其中这只是一个简单的示意图，实际上每一层都可以写入指纹向量。
圆形指纹算法：
- 输入：
  - 分子结构
  - 半径参数
  - 指纹向量长度
- 输出：指纹向量
- 算法步骤：
  - 初始化指纹向量：
  - 遍历每个原子，获取每个原子的特征
  - 遍历每一层。对于第层，迭代步骤为：
    - 遍历分子中的每个原子，对原子计算：
      - 获取顶点邻域原子的特征：
      - 拼接顶点及其邻域原子特征：
      - 执行哈希函数得到顶点的当前特征：
      - 执行索引函数：
      - 登记索引：
  - 最终返回

5.1.2 分子指纹GCN算法

我们选择类似于现有ECFP 的神经网络架构：
- 哈希操作Hashing：在ECFP 算法中，每一层采用哈希操作的目的是为了组合关于每个原子及其邻域子结构的信息。
  
  我们利用一层神经网络代替哈希运算。当分子的局部结构发生微小的变化时（神经网络是可微的，因此也是平滑的），这种平滑函数可以得到相似的激活值。
- 索引操作Indexing：在 ECFP 算法中，每一层采用索引操作的目的是将每个原子的特征向量组合成整个分子指纹。每个原子在其特征向量的哈希值确定的索引处，将指纹向量的单个比特位设置为1，每个原子对应一个1 。这种操作类似于池化，它可以将任意大小的Graph 转换为固定大小的向量。
  
  这种索引操作的一个缺点是：当分子图比较小而指纹长度很大时，最终得到的指纹向量非常稀疏。
  
  我们使用softmax 操作视作索引操作的一个可导的近似。本质上这是要求将每个原子划分到一组类别的某个类别中。所有原子的这些类别向量的总和得到最终的指纹向量。其操作也类似于卷积神经网络中的池化操作。
- 规范化Canonicalization：无论原子的邻域原子的顺序如何变化，圆形指纹是不变的。实现这种不变性的一种方式是：在算法过程中，根据相邻原子的特征和键特征对相邻原子进行排序。
  
  我们尝试了这种排序方案，还对局部邻域的所有可能排列应用了局部特征变换。
  
  另外，一种替代方案是应用排序不变函数permutation-invariant，如求和。为了简单和可扩展性，我们选择直接求和。
神经网络指纹算法：
- 输入：
  - 分子结构
  - 半径参数
  - 指纹长度
  - 隐层参数，输出层参数。
    
    对不同的键数量，采用不同的隐层参数（最多五个键）。
- 输出：指纹向量
- 算法步骤：
  - 初始化指纹向量：
  - 遍历每个原子，获取每个原子的特征
  - 遍历每一层。对于第层，迭代步骤为：
    - 遍历分子中的每个原子，对原子计算：
      - 获取顶点邻域原子的特征：
      - 池化顶点及其邻域的原子的特征：
      - 执行哈希函数：，为邻域顶点数量。
      - 执行索引函数：
      - 登记索引：
  - 最终返回
设指纹向量的长度为，顶点特征向量的维度为，则的参数数量为，的参数数量为。
上述 ECFP 算法和神经网络指纹算法将每一层计算得到的指纹叠加到全局指纹向量中。我们也可以针对每一层计算得到一个层级指纹向量，然后将它们进行拼接，而不是相加。

以神经网络指纹算法为例：
- 在第层计算索引为：。然后登记索引：
- 最终将所有层的索引拼接：
ECFP 圆形指纹可以解释为具有较大随机权重的神经网络指纹算法的特殊情况。

在较大的输入权重情况下，当为时，该激活函数接近阶跃函数。而级联的阶跃函数类似于哈希函数。

在较大的输入权重情况下，softmax 函数接近一个one-hot 的 argmax 操作，这类似于索引操作。

5.1.3 限制

计算代价：神经网络指纹在原子数、网络深度方面与圆形指纹具有相同的渐进复杂度，但是由于在每一步都需要通过矩阵乘法来执行特征变换，因此还有附加的计算复杂度。

假设分子的特征向量维度为，指纹向量长度为，网络深度为，原子数量为，则神经网络指纹的计算复杂度为。

实际上在圆形指纹上训练一个单隐层的神经网络只需要几分钟，而对神经网络指纹以及指纹顶部的单隐层神经网络需要一个小时左右。
每层的计算限制：从网络的一层到下一层之间应该采取什么结构？本文采用最简单的单层神经网络，实际上也可以采用多层网络或者 LSTM 结构，这些复杂的结构可能效果更好。即：用复杂的神经网络”block”结构代替简单的神经网络单层结构。
图上信息传播的限制：图上信息传播的能力受到网络深度的限制。对于一些规模较小的图如小分子的图，这可能没有问题；对于一些大分子图，这可能受到限制。

最坏情况下，可能需要深度为的网络来处理规模（以原子数来衡量）为的图。

为了缓解该问题，Spectral networks and locally connected networks on graphs 提出了层次聚类，它只需要层就可以在图上传播信息。这种方式需要解析分子为树结构，可以参考NLP 领域的相关技术。
无法区分立体异构体stereoisomers ：神经网络指纹需要特殊处理来区分立体异构体，包括enantomers 对映异构体（分子的镜像）、cis/trans isomers 顺/反异构体（绕双键旋转）。大多数圆形指纹的实现方案都可以区分这些异构体。

5.2 实验

5.2.1 随机权重

分子指纹的一个用途是计算分子之间的距离。这里我们检查基于 ECFP 的分子距离是否类似于基于随机的神经网络指纹的分子距离。

我们选择指纹向量的长度为 2048，并使用Jaccard 相似度来计算两个分子的指纹向量之间的距离：

我们的数据集为溶解度数据集，下图为使用圆形指纹和神经网络指纹的成对距离散点图，其相关系数为。

图中每个点代表：相同的一对分子，采用圆形指纹计算到的分子距离、采用神经网络指纹计算得到的分子距离，其中神经网络指纹模型采用大的随机权重。

距离为1.0 代表两个分子的圆形指纹没有任何重叠；距离为0.0 代表两个分子的圆形指纹完全重叠。
我们将圆形指纹、随机神经网络指纹接入一个线性回归层，从而比较二者的预测性能。
- 圆形指纹、大的随机权重的随机神经网络指纹，二者的曲线都有类似的轨迹。这表明：通过大的随机权重初始化的随机神经网络指纹和圆形指纹类似。
- 较小随机权重初始化的随机神经网络指纹，其曲线与前两者不同，并且性能更好。
- 即使是未经训练的神经网络，神经网络激活值的平滑性也能够有助于模型的泛化。

5.2.2 可解释性

圆形指纹向量的特征（即某一组bit 的组合）只能够通过单层的单个片段激活（偶然发生的哈希碰撞除外），神经网络指纹向量的特征可以通过相同结构的不同变种来激活，从而更加简洁和可解释。
为证明神经网络指纹是可接受的，我们展示了激活指纹向量中每个特征对应的亚结构类别。
- 溶解性特征：我们将神经网络指纹模型作为预训溶解度的线性模型的输入来一起训练。下图展示了对应的片段（蓝色），这些片段可以最大程度的激活神经网络指纹向量中最有预测能力的特征。
  - 上半图：激活的指纹向量的特征与溶解性具有正向的预测关系，这些特征大多数被包含亲水性R-OH 基团（溶解度的标准指标）的片段激活。
  - 下半图：激活的指纹向量的特征与溶解性具有负向的预测关系（即：不溶解性），这些特征大多数被非极性的重复环结构激活。
- 毒性特征：我们用相同的架构来预测分子毒性。下图展示了对应的片段（红色），这些片段可以最大程度的激活神经网络指纹向量中最有预测能力的特征。
  - 上半图：激活的指纹向量的特征与毒性具有正向的预测关系，这些特征大多数被包含芳环相连的硫原子基团的片段激活。
  - 下半图：激活的指纹向量的特征与毒性具有正向的预测关系，这些特征大多数被稠合的芳环（也被称作多环芳烃，一种著名的致癌物）激活。

5.2.3 模型比较

数据集：我们在多个数据集上比较圆形指纹和神经网络指纹的性能：
- 溶解度数据集：包含 1144 个分子，及其溶解度标记。
- 药物功效数据集：包含 10000 个分子，及其对恶行疟原虫（一种引发疟疾的寄生虫）的功效。
- 有机光伏效率数据集：哈佛清洁能源项目使用昂贵的 DFT 模拟来估算有机分子的光伏效率，我们从该数据集中使用 20000 个分子作为数据集。
我们的 pipeline 将每个分子编码的 SMILES 字符串作为输入，然后使用 RDKit 将其转换为Graph 。我们也使用 RDKit 生成的扩展圆形指纹作为 baseline 。这个过程中，氢原子被隐式处理。

我们的 ECFP 和神经网络中用到的特征包括：
- 原子特征：原子元素类型的 one-hot、原子的度degree、连接氢原子的数量、隐含价implicit valence、极性指示aromaticity indicator。
- 键特征：是否单键、是否双键、是否三键、是否芳族键、键是否共轭、键是否为环的一部分。
我们采用 Adam 优化算法，训练步数为 10000，batch size = 100 。我们还使用了 batch normalization 技术。
- 我们还对神经网络进行了 tanh 和 relu 激活函数的对比实验，我们发现relu 在验证集上一直保持优势并且优势不大。
- 我们还对神经网络进行了 drop-connect 实验，它是 dropout 的一个变种，其中权重被随机设置为零（而不是隐单元被随机设置为零）。我们发现这会导致更差的验证误差。
我们使用 Random-Search 来优化以下超参数：学习率的对数、初始权重的对数、正则化系数的对数、指纹向量长度、指纹深度（最深六层）、全连接网络层的维度、神经网络指纹的隐层维度。

所有超参数使用折交叉验证来优化，其中每一折随机执行50 次。
我们比较了两种情况下圆形指纹和神经网络指纹的性能：
- 第一种情况：一个线性层使用指纹向量作为输入来执行预测，即 linear layer。
- 第二种情况：一个单隐层的神经网络使用指纹向量作为输入来执行预测，即 neural net 。
结果如下图所示。可以看到在所有实验中，神经网络指纹均达到或者超过圆形指纹的性能，并且使用神经网络层的方式（neural net ）超过了线性层的方式（linear layer）。