[论文笔记] [2016] Molecular graph convolutions: moving beyond fingerprints

这篇论文提出了一种GNN变体应用于 virtual screening 中分子表征学习阶段，取代了传统用分子指纹做分子表征。这也是目前分子表征主流的一种方式。
本文中图卷积模块 Weave module，如下图，其中了考虑边的特征，并且提出了一种新的图表示方法，即一种新的聚合节点信息得到全局表达的方式。

Methods

Desired invariants of a model

设计一个DL结构作用于分子图上，作者认为需要满足三个property。

第一个是，The output of the model should be invariant to the order that the atom and bond information is encoded in the input. 这点即是要求输入分子图，不管原子和化学键的顺序是怎么样的，经过图读出操作（ReadOut）得到的分子表征都是一样的。这里的核心要义在于：操作本身要对节点顺序不敏感。
这里就有一个图重构（Graph Isomorphism）问题，即对一个图上节点重新编号，形成的图与原先的图其实是等价的。为了使得同构图的表示能够保持一致，图读出的操作就需要对结点顺序不敏感。在数学上，能表达这种操作的函数也称为对称函数。

在模型的结构中，存在着两种表征，一个 atom layer，包含与分子图上原子相关的二维矩阵，用 $A_a^x$ 表示 atom layer x 的原子a的向量（特征），一个是 pair layer，包含着与分子图上原子对（也就是边）相关的三维矩阵，用$P_{(a,b)}^x$ 表示 pair layer 中 pair (a, b) 的向量（特征）。

第二个 property 是，（Atom and pair permuation invariance）The values of an atom layer and pair permute with the original input layer order. More precisely, if the inputs are permuted with a permutation operator Q, then for all layers x, y, $A^x$ and $P^y$ are permuted with operator Q as well. 这个 property 意味着从atom layer（或者 pair layer）中单独的原子来看，它的向量不会随着原子顺序变化而变化。

第三个 property 是，（Pair order invariance）For all pair layer y, $P_{(a,b)}^y=P_{(b,a)}^y$，因为分子图是无向图，即边（a, b）和边（b, a）的向量（特征）是要相等的，那么就需要经过操作后满足上述 property。

在实现满足第二个和第三个property的操作后，就能由atom和pair的特征信息构造出分子层面的表征，并且满足property 1。

Invariant-preserving operations

作者首先定义了两个函数 $f$ 和 $g$，$f$ 是一个任意函数，$g$是一个任意的 commutative function（这里就是上面提到的对称函数），对参数的顺序，也就是对节点顺序并不敏感。在这篇论文中，$f$是一个带ReLU的线性变换函数（成了非线性），$g$ 是一个求和函数。

根据上面的两个函数，作者定义了一系列包含在Weave module中的操作：
($A\to A$) 与 ($P\to P$)
这两个操作，是作用在atom（pair）layer上的，其定义如下：
$$A_a^y=f(A_a^{x1},A_a^{x2},...,A_a^{xn})$$
$$P_{a,b}^y=f(P_{a,b}^{x1},P_{a,b}^{x2},...,P_{a,b}^{xn})$$
上面两个操作的作用是聚合每个atom（pair）所对应的单层或多层的特征，生成为新一层的atom特征，并且维持了property2。这边$(..,...)$代表着特征的拼接。值得注意的是，这里的操作函数作用于每个atom是同一个，那么意味着是参数共享。于是上述两个操作可以近似看成图上的卷积操作。

($A\to P$) 和 ($P\to A$)
这两个操作是weave module的核心，即定义了GCN如何做消息传递。之前看到的GCN都是只有atom间的消息传递，并没有包含边的信息，而本文包含了边的特征，那如何做边的信息传递则需要去考虑。作者提出的两个操作，(P $\to$ A) 实现了atom的信息传递，(A $\to$ P)实现了pair的信息传递。
atom的信息传递操作如下：
$$A_a^y=g(f(P_{a,b}^x),f(P_{a,c}^x),f(P_{a,d}^x),...)$$
上式可以看出，传递到atom a的消息聚合方式是对连接atom a的各个pair的特征经过非线性变换后的求和。
而对于pair的信息传递操作如下：
$$P_{ab}^y=g(f(A_a^x,A_b^x),f(A_b^x,A_a^x))$$
不难发现，这个操作是能维持property3的，即$P_{ab}^y$ 与 $P_{ba}^y$的结果是一样的。

将上面的四个操作结合，就能构建一个weave module，也就是最上面那个结构，可以看成是一个图卷积层。但作者并没有对这个结构做太多解释，让人看了会有点迷，为什么会这么设计。但如果将这个模型带入到MPNN框架[2]中去理解，会清楚很多。

MPNN框架中消息传递阶段，有两个函数$M_t$函数和$U_t$，以中间状态$h_v^t$（在这里就是$A_a^x$），消息$m_v^{t+1}$和边的特征向量$e_{vw}^t$（这里就是pair的特征向量$P_{ab}^x$）。

如果对应到weave module中，消息函数为$M_k(A_a^k,A_b^k,P_{ab}^k)=\sigma (W_0{P}_{ab}^k)$，那么原子a聚合邻居消息得到的$m_a^{k+1}={\sum }_{w\in N(v)}{M}_k(h_v^t,h_w^t,e_{vw})={\sum }_{i\in N(v)}\sigma (W_0{P}_{ai}^k)$，其中$\sigma$是ReLU，$W_0$即上面$f$函数的可学习参数。可以看出求$m_a^{k+1}$的过程就是上面结构中的$(P\to A{)}_0$操作，得到了$A_a^{k^{\prime \prime }}$。而更新函数$U_k(A_a^k,m_a^{k+1})=\sigma (W_2(\sigma (W_1{A}_a^k),m_a^{k+1}))$，其中$\sigma (W_1{A}_a^k)$ 就是上面结构中的 $(A\to A{)}_0$操作，得到了 $A_a^{k^{\prime }}$，而$\sigma (W_2(...,...))$便是上面结构中$(A\to A{)}_1$操作，即得到 $A_a^{k+1}$，atom a的特征完成更新。另一边，边的特征更新为 $P_{ab}^{k+1}=U_t^{\prime }(P_{ab}^k,A_a^k,A_b^k)=\sigma (W_5(\sigma (W_3{P}_{ab}^k),\sigma (W_4(A_a^k,A_b^k))+\sigma (W_4(A_b^k,A_a^k)))))$，其中$\sigma (W_3{P}_{ab}^k)$即为上面的$(P\to P{)}_0$操作，而$\sigma (W_4(A_a^k,A_b^k))+\sigma (W_4(A_b^k,A_a^k))$则为$(A\to P{)}_0$操作，$\sigma (W_5(...,...))$为$(P\to P{)}_1$操作。

Molecule-level feature

这篇论文的另外一个亮点便在于根据原子特征和边特征生成分子层面的特征，也就是实现图读出操作。读出操作与CNN模型中常用的紧跟最后一个卷积层的全局池化（Global Pooling）操作如出一辙。两者都是通过对所有输入的一次性聚合得到全局表达，且与池化的常见类型一样，读出机制也可以取sum、mean、max等类型的函数。

sum、mean、max这些方式很容易实现的，但一个问题就是信息丢失太多，本来一张图有n个节点，每个结点m维，而一经压缩，得到的只是一个m维的向量。作者提出了一种用类似直方图的方法来对每维数据分布进行建模，[1]中作者举了一个例子展示了直方图能巧妙的平衡数据的压缩与增强。利用直方图来压缩数据的大体流程是：把某一维上的特征值的取值范围划分为n个区间（即桶），根据特征值所在区间，将特征值装入桶中，最后对每一个桶内数值求和，这样便得到n维的特征。而作者此处并没有直接采用直方图方法，而是采用由高斯核实现的模糊直方图（Fuzzy histograms）。其数据压缩过程为：将某一维上的特征值的取值范围内划分为n个区间，而区间的边界点为n个高斯分布的均值，方差预先定义好；这样对于某一维的某个特征值，其与每个高斯分布交点的纵坐标为落入该分布中的取值，然后将所有数值归一化，就得到了该特征值分布在各个区间的比例。如下图，一个特征值为1.8，其与3个高斯分布的交点纵坐标分别为0,0.3,0.9，将其归一化，得到一个3维的向量[0.0, 0.25, 0.75]。

利用上述方法，就能得到一个信息损失较少的分子层表示。另外，在具体实现中，有一个细节是作者对图读出操作前，作者最后用了一层作用于每个原子的卷积，即$(A\to A)$操作，将每个原子的特征维度进行扩展，意图将原先维度信息分散到更广的维度（降低每个特征的重要性），这样在后面进行图读出操作，压缩数据到一个低维的特征时，就能减少信息的损失。这倒是一个值得借鉴的trick。

另外，在实验过程中，作者在初始化原子特征和化学键特征时，设置了一组实验，采用simple feature（原子特征只有原子类型，化学键特征只有化学键类型与graph distance），与全量特征的效果，效果差距很少。

总结

这篇论文的贡献在于将GNN应用于分子表示，打开了分子表示的一个全新的思路即分子表示学习，还有一个亮点就是图读出操作，另外便是一些trick和将GNN应用于分子表示学习的具体实验细节。

参考文献

[1] 从图(Graph)到图卷积(Graph Convolution): 漫谈图神经网络 (三), https://zhuanlan.zhihu.com/p/108299847
[2] Gilmer J, Schoenholz S S, Riley P F, et al. Neural message passing for quantum chemistry[C]//Proceedings of the 34th International Conference on Machine Learning-Volume 70. JMLR. org, 2017: 1263-1272.