阅读笔记-KDD2020-MoFlow: An Invertible Flow Model for Generating Molecular Graphs

相关资料：

• Slides
• Video
• Paper
• Code

一、 Introduction

这篇文章是$flow-based$，大致就是有两个概率分布，如果他们之间的变换是可逆的，那么就可以根据一个较为简单地分布通过逆变换获得另外一个较复杂的概率分布。这篇文章采用的就是将邻接矩阵和原子矩阵这两个概率分布与另外两个较为简单的潜在变量的分布（文中将这两个潜在变量的分布采样为高斯分布）利用flow联系了起来。通过设计一个可逆的分布，以及从高斯分布中采样，将采样的潜在变量通过这个可逆的变换就可以得到分子图结构。
生成过程如下：

二、Inference（encoding）

2.1 Model Preliminary.

2.1.1 The flow framework.

给定两个分布：

• $\mathcal{Z}\sim {\mathbf{P}}_{\mathcal{Z}}(\mathbf{Z})$：较为容易建模的一个概率分布（在这样一个潜在的空间里，有很强的独立性）；
• $\mathcal{X}\sim {\mathbf{P}}_{\mathcal{X}}(\mathbf{X})$：较为复杂的高维数据；

$flow-based$的模型旨在找到一系列的可逆变换${\mathbf{f}}_{\mathbf{\Theta }}={\mathbf{f}}_{\mathbf{L}}\circ ...\circ {\mathbf{f}}_1$，来实现这两个概率分布之间的变换。并且当$\mathbf{Z}={\mathbf{f}}_{\mathbf{\Theta }}(\mathbf{X})={\mathbf{f}}_{\mathbf{L}}\circ ...\circ {\mathbf{f}}_1(\mathbf{X})$时，有${\mathbf{P}}_{\mathcal{X}}(\mathbf{X})={\mathbf{P}}_{\mathcal{Z}}(\mathbf{Z})|\mathbf{d}\mathbf{e}\mathbf{t}(\frac{\partial \mathbf{Z}}{\partial \mathbf{X}})|$。所以有：

2.1.2 Invertible affine coupling layers.

然而，如何设计一个满足以下三个条件的可逆变换：（a.）这个变换是一个可逆的函数$f_{\Theta }$；（b.）这个可逆变换富有表现力；（c.）这个可逆变换可以高效计算雅可比矩阵的行列式。
作者设计了如下的变换来满足上述条件：

• $\mathbf{t}\mathbf{h}\mathbf{e}\mathbf{S}\mathbf{c}\mathbf{a}\mathbf{l}\mathbf{e}\mathbf{f}\mathbf{u}\mathbf{n}\mathbf{c}\mathbf{t}\mathbf{i}\mathbf{o}\mathbf{n}{S}_{\Theta }:{\mathbb{R}}^d\mapsto {\mathbb{R}}^{n-d}$
• $\mathbf{t}\mathbf{h}\mathbf{e}\mathbf{T}\mathbf{r}\mathbf{a}\mathbf{n}\mathbf{s}\mathbf{f}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{a}\mathbf{t}\mathbf{i}\mathbf{o}\mathbf{n}\mathbf{f}\mathbf{u}\mathbf{n}\mathbf{c}\mathbf{t}\mathbf{i}\mathbf{o}\mathbf{n}{T}_{\Theta }:{\mathbb{R}}^d\mapsto {\mathbb{R}}^{n-d}$
• $det(\frac{\partial Z}{\partial X})=exp({\Sigma }_j{S}_{\Theta }(X_{1:d}{)}_j).$

2.1.3 Splitting Dimensions.

压缩空间维度将$X^{c\times n\times n}$转换为$X^{(c{h}^2)\times \frac{n}{h}\times \frac{n}{h}}$来获得更多的channels并将这些channels分隔开，分别送进模型中。

• 2.1.3 Splitting Dimensions.

2.1.4 Numerical stability by actnorm

为了保证基于流的模型的数值稳定性，Glow中引入了actnorm层，该层通过一个具有可学习尺度和偏差的仿射变换对批次内每个通道的尺寸进行规范化。刻度和偏差被初始化为Batch中每个通道的平均值和标准差的倒数。

2.2 Proposed Moflow Model

2.2.1 Problem Definition: Learning a ProbabilityModel of Molecular Graphs

• $\mathcal{A}\in {\mathcal{R}}^{n\times k}$：原子矩阵，并且这个原子矩阵最多具有$n$个，以及$k$种类型的原子。$\mathcal{A}(i,k)=1$代表着原子$i$的类型为$k$；
• $\mathcal{B}\in {\mathcal{R}}^{c\times n\times n}$：化学键矩阵。$c$：代表着化学键类型的数量。$B(c,j,i)=1$代表着原子$i$与原子$j$之间的化学键类型为$c$；
• $\mathcal{M}=\mathcal{A}\times \mathcal{B}\in {\mathbb{R}}^{n\times k}\times {\mathbb{R}}^{c\times n\times n}$，即矩阵$\mathcal{M}$是矩阵$\mathcal{A}$和$\mathcal{B}$的笛卡尔积。因此分子图结构$\mathcal{M}$可以看做具有多种类型的节点与多种类型的边的无向图。

为了捕捉分子图中原子和键的组合结构，作者将这种分子生成模型分解为两个部分：${\mathbf{P}}_{\mathcal{M}}(\mathbf{M})={\mathbf{P}}_{\mathcal{M}}((\mathbf{A},\mathbf{B}))\approx {\mathbf{P}}_{\mathcal{A}|\mathcal{B}}(\mathbf{A}|\mathbf{B};{\theta }_{\mathcal{A}|\mathcal{B}}){\mathbf{P}}_{\mathcal{B}}(\mathbf{B};{\theta }_{\mathcal{B}})$

参数${\theta }_{\mathcal{B}}与{\theta }_{\mathcal{A}|\mathcal{B}}$通过最大化下式来学习到：

模型包含两个部分：

• $graphconditionalflowforatoms$；
• $flowforbonds$；

作者还提出了进一步学习了学习到的潜在向量和分子特性之间的映射关系，以回归基于图的分子特性，并指导优化分子图的生成。

2.2.2 Glow for Bonds

$Z_B=f_{\mathcal{B}}(B)$，并且$Z_B$满足多元高斯分布，由于变换$f$是可逆的，所以可得：

我们先从$P_{\mathcal{Z}}(Z)$中采样$\tilde{Z}$，并根据$f$的逆变换得到边矩阵即$\tilde{B}=f_{\mathcal{B}}^{-1}(\tilde{Z})$

具体过程如下图所示：

2.2.3 Graph Conditional Flow for Atoms

• 给定化学键矩阵$B\in \mathcal{B}\subset {\mathbb{R}}^{c\times n\times n}$，$atomflow$的目标是为了生成正确的原子矩阵$A\in \mathcal{A}\subset {\mathbb{R}}^{n\times k}$从而来组装有效的分子$M=(A,B)\in \mathcal{M}\subset {\mathbb{R}}^{n\times k+c\times n\times n}$。作者先定义了$\mathcal{B}$-conditional flow and graph conditional flow$f_{\mathcal{A}|\mathcal{B}}$。这个变换是根据给定的矩阵$\mathcal{B}$，来将矩阵$\mathcal{A}$转换为${\mathcal{Z}}_{A|B}$即$Z_{A|B}=f_{\mathcal{A}|\mathcal{B}}(\mathcal{A}|\mathcal{B})$，并且$P_{{\mathcal{Z}}_{\mathcal{A}|\mathcal{B}}}$是一个 各向同性的多元高斯分布。

• $Z_{A|B}=f_{\mathcal{A}|\mathcal{B}}(\mathcal{A}|\mathcal{B})$是一个可逆的变换（即存在$f_{\mathcal{A}|\mathcal{B}}^{-1}(Z_{A|B}|B)=A|B$），并且变换前后维数保持不变。接下来来求取满足flow模型的雅克比矩阵。因为$B\in \mathcal{B}$在整个变换过程中保持不变。所以有
  
  则有
  则有:
  

• $\mathbf{g}\mathbf{r}\mathbf{a}\mathbf{p}\mathbf{h}\mathbf{c}\mathbf{o}\mathbf{u}\mathbf{p}\mathbf{l}\mathbf{i}\mathbf{n}\mathbf{g}\mathbf{l}\mathbf{a}\mathbf{y}\mathbf{e}\mathbf{r}$：
  
  则：
  

• actnorm2D：来执行一个正则化
  
  则有
  接下来作者介绍了如何使得网络变得更深的一个技巧：即每隔一层交换$A_1,A_2或者B_1,B_2$的位置。如下所示：

2.2.4 Validity Correction

分子必须遵循每个原子的化合价约束，但是用生成的键张量和原子矩阵组装一个分子可能会导致化学无效。所以需要为每个原子定义一个有效性检查：

作者举例说明了这一形式。对于铵根正离子$[N{H}_4{]}^{+}$中的$N$它会有四个化学键，而不是三个。同样的，$S^{+}与O^{+}$也会有3个化学键而不是两个。作者将$O^{+},S^{+},N^{+}$的$Ch$设为1，其他原子的$Ch$设为0。

作者设计的有效性检查尝试对分子做最小的改变来使其具有化学上的合理性，并使分子图中最大的连接成分尽可能大。有效性检查如下：

• 首先，根据上式，为每个原子检查化学键数目的有效性；
• 如果分子满足上述约束，将图结构中最大的连接成分作为最终生成的分子输出，整个生成过程结束；
• 如果分子中第$i$个原子没有满足上述约束，则对按顺序对这个原子的化学键进行排序，并将顺序最大的键减去1，将得到的新分子重新进行第一步。

三、Experiments

作者通过以下四个方面进行了实验来评估MoFlow模型：

• 分子图生成与重建：首先是重建数据集中的所有分子，并生成尽可能多的有效的，不属于数据集的分子图；
• 可视化潜在空间：MoFlow能否将分子图嵌入到具有合理化学相似性的连续潜空间中?
• 性质优化：模型能否生成具有优化性质的分子图；
• 约束属性优化：，模型能否生成具有优化特性的新颖分子图，同时尽可能保持化学相似性？

作者采用的数据集如下：

作者在文章中提到，对于$ZINC250K$MoFlow只用了22小时便训练完成，而基于自回归的方法$GraphAF$用了124个小时，这也体现了非自回归方法高效的优点。

3.1 生成与重建

五个指标：

• Validity：化学有效分子在所有生成的分子中的百分比；
• Uniqueness：所有生成的分子中不属于原数据集并且满足化学有效性的分子数与生成的分子数之比；
• Novelty：所有生成的分子中不属于原数据集并且满足化学有效性的分子数与数据集所含分子数之比（貌似和上一个指标有些重复，可能上一个指标将生成的分子中重复的分子删去后在进行计算）；
• Reconstruction：原数据集中能够完全被模型重建的分子数与数据集所含分子数之比；
• N.U.V.：生成的分子中满足：Validity，Uniqueness，Novelty分子的数量与生成的分子数之比。
• Validity without check/correction

实验结果如下：

MoFlow在不进行化学有效性检查的情况下，仍然获得了极佳的生成结果。

3.2潜在空间的可视化

作者利用谷本系数来反应化学相似性。
作者先讲一个分子映射为一个潜在向量Z，通过下式来得到与之相近的分子：

X,Y均为单位矢量。
结果如下：

3.3 性能优化

这个任务的目的是为了生成具有最佳QED分数的分子。
作者先将分子图映射为嵌入向量$Z$，然后训练一个多层的MLP来回归预测其QED分数。训练完成后，根据先将Z送入MLP中得到QED分数y，然后根据梯度来优化$Z$即：

执行K次后，得到最终的优化分子，经过逆变换生成对应的分子。实验结果表明，MoFlow模型能够找到更多具有最佳QED分数的分子图结构：

3.4 约束属性优化

这个任务致力于给定一个分子图，来找到一个符合两个最值的分子，即与原分子最相似，在相关性能分数上得到最大的提升。
作者使用了morgan指纹的tanimoto相似性系数来衡量两个分子之间的相似性，以及$logP$作为目标性质。作者同样用基于梯度的方法来实现。实验结果如下：