今天给大家介绍的是来自德州农工大学的Shuiwang Ji等人发表在预印网站arXiv的文章GraphDF: A Discrete Flow Model for Molecular Graph Generation。在分子图的生成问题中。虽然分子图是离散的，但大多数现有的方法都使用连续的隐变量，从而导致对离散的图结构不正确的建模。在这项工作中，作者提出了一种新的基于流并使用离散隐变量的分子图生成模型。GraphDF使用可逆模移位变换，将离散的隐变量映射到图节点和边。作者表示，使用离散的隐变量降低了计算成本，并消除了反离散化的负面影响。实验结果表明，GraphDF在随机生成、性质优化和约束优化任务上的性能优于以往的方法。
一、研究背景
药物发现与化学领域的一个基本问题是设计具有特定性质的新分子。这个任务目前仍然具有挑战性，因为分子空间自离散且巨大。在此领域中深度生成模型已经取得了显著的进展。许多最近的研究将分子结构表示为图，并提出用先进的深度生成模型生成新分子图。这些方法首先将分子图映射到连续隐空间中的向量。在生成时，生成模型会从隐空间中随机选择一个连续向量，并映射回分子图空间。
为了使用具有连续隐变量的生成模型，现有方法通过添加实值噪声将离散图数据转化为连续数据。然而，这种反离散化处理阻止了模型从离散图结构中获得原始离散分布，从而增加了模型训练的难度。这使得模型很难获得图结构的真实分布并产生多样的分子。
在本文中，作者提出了GraphDF，一个使用离散隐变量生成分子图的生成模型。GraphDF通过序列化采样离散隐变量来生成分子图，并通过可逆的模移位变换将它们映射到新的节点和边。与以往的方法相比，作者提出的离散变换消除了计算雅可比矩阵的成本，克服了反离散化所造成的限制。实验结果表明，GraphDF在各种分子生成任务上都优于以往最先进的方法。
二、模型与方法
2.1 序列生成的概述
设k为节点类型的数量，c为边类型的数量。作者使用G=(X,A)来表示具有n个节点的分子图，其中X∈{0,1}n×k为节点矩阵，A∈{0,1}n×n×c为邻接张量。X[i,u]=1表示第i个节点为类型u，A[i,j,v]=1表示在第i个节点和第j个节点中间存在一个类型为v的边。如果第i个节点和第j个节点中间不存在边，那么A[i,j,v]=0，v=0,…,c-1。
类似于MolecularRNN，作者通过序列决策过程生成分子图。从一个空的分子图G0开始。在步骤i中，当前的子图Gi-1包含i-1个节点。生成模型fθ首先基于采样的隐变量zi生成一个类型为ai的新节点，其中ai属于{0,…,k-1}。然后，由fθ按顺序生成新节点与Gi-1中所有节点之间的边。具体地说，Gi-1的第j个节点的边是基于隐变量zij生成的，其边类型bij∈{0，…，c-1,c}，其中额外的边类型c表示没有边。如果bij违反了化学键的价键规则，则会重新采样一个新的zij来重新生成它。上述过程是首先生成元素，然后重复直到在步骤i没有添加新边的自回归函数。这个过程的一个例子如图1所示。总之，生成分子图G相当于通过下式生成序列表示SG=(a1,a2,b21,a3,b31,b32,a4,…)。
图1. 序列生成过程的图示
2.2 使用离散的隐变量进行生成
在作者的方法中，所有的隐变量都是离散的，并从多项式分布中采样。具体地说，zi是从一个具有可训练参数(α0,…,αk-1)的多项式分布pZa中采样的：
类似地，zij也可以从具有可训练参数(β0,…,βc)的多项式分布pZb中采样。我们使用一个离散的流模型来可逆地将离散的隐变量映射到新的节点和边。离散流中使用的离散变换是一种模移位变换，形式为：
其中t是类别数，z，μ∈{0,…,t-1}。形式上，用于生成新节点和边的离散变换由D个模移位模块组成：
其中f o g(·)被定义为f(g(·))。
这里，移位因子µdi和µdij分别是(a1,a2,b21,…,bi-1,i-2)和(a1,a2,b21,…,bi,j-1)的函数。它们都扮演着从以前生成的元素中捕获条件信息的作用。图2中给出了一个用该框架生成简单分子图的例子。
图2. 使用GraphDF框架生成分子图的例子
使用离散隐变量是该方法与其他基于流模型的生成方法的关键区别。此模型有几个优点。离散可逆映射的公式不涉及雅可比矩阵。因此，采样节点和边的概率可以直接得到：
这导致了计算成本的大幅降低。
更重要的是，离散的隐变量可以使生成模型免于反离散化的缺陷。以前的方法通过向独热编码中添加实值噪声来将图节点和边的类别表示转换为连续数据。然后连续流模型将这种伪连续数据映射到服从于高斯分布的隐变量。作者认为这种反离散化操作可能是有问题的。首先，生成模型并没有捕获到原始离散数据的真实分布，而是捕获了反离散化后的失真分布。其次，由于噪声的随机性，即使对同一个图执行多次去量化，所得到的连续数据也可能是非常不同的。因此，在训练过程中，单个图可能会映射到隐空间中的两个相距很远的点，导致模型振荡与不收敛。然而，GraphDF将离散图数据映射到离散隐空间，不引入任何分布失真。因此，它可以更精确地建模图结构的底层分布，从而消除了反离散化的影响。
2.3 推断与训练
生成的反过程是推断，它包括从图映射到离散的隐变量。给定一个图G=(X,A)，首先根据广度优先遍历(BFS)顺序对它的节点和边重新排序来计算它的序列表示SG。然后，可以通过下式得到该序列中每个元素对应的离散隐变量。
其中
与生成过程类似，移位因子µdi和µdij分别在ai和bij之前从SG中的元素获取自回归信息。
作者通过最大化所有数据样本的对数似然来训练GraphDF模型。给定一个分子图G，首先对G进行推断来计算对应于它的节点和边的隐变量。然后G的对数似然可以计算为：
GraphDF模型和离散先验分布pZa，pZb的参数通过梯度下降更新，以最大限度地提高所有数据样本的对数似然。
2.4 不变性和离散条件生成
生成新节点和边的离散变换是自回归的，需要将条件信息捕获到µdi和µdij中。Tran等人使用语言模型，比如LSTM和Transformer，在离散流模型中开发序列生成器。然而，相同的方法不能直接用于生成SG，原因有两个，如下：
·语言模型只了解SG中编码的节点和边的类型信息，而不捕获图中丰富的结构信息。
·在中间阶段生成的子图并没有内在的顺序，但语言模型在处理SG时假设了一个顺序。例如，对于（a1,a2)=（1,0）和（a1,a2)=（0,1），生成模型应给出相同的估计p（b21|a1,a2），但语言模型不能提供这样的保证。
为了解决这两个问题，作者提出利用图卷积网络(GCN)从中间子图中获取条件信息。具体地说，设在添加节点ai之前生成的子图为Gi-1=(X,A)，作者使用L层关系GCN(RGCN)通过执行L消息传递操作来从Gi-1获得节点嵌入HL∈Rn×r。第l条消息传递被如下执行：
其中，
，Dv为对角矩阵，
，{Wlv}cv=1为第l层的训练权重矩阵。注意，在这个离散流中所有模移位模块之间共享相同的R-GCN。然后，利用基于多层感知器(MLP)的分类网络来计算µdi：
其中h∈Rr为在节点嵌入HL上执行sum-池化获得的图嵌入，MLPda将输入投影到一个k维对数向量，并且argmax输出对数值最大的位置索引。由于argmax是不可微的，作者在训练过程中使用softmax温度函数来近似梯度：
其中，作者在实验中固定了温度τ=0.1。
类似地，为了获得µdij，作者使用R-GCN从添加边bij之前生成的子图Gi-1,j中提取节点嵌入HL和图嵌入h。然后作者用MLP计算µdij：
其中，HLi和HLj分别表示Gi-1,j中第i个和第j个节点的节点嵌入，Con(·)是连接操作。softmax近似也在这里通过argmax被用于执行后向传播。
2.5 使用强化学习的性能优化
上述讨论只涉及到了与性质分数无关的分子图生成。然而，性质优化和约束优化都要求生成的分子具有优化的性质分数。为此，我们使用强化学习来微调一个训练好的GraphDF模型来优化指定的性质分数。
作者将使用GraphDF的序列生成过程表示为马尔可夫决策过程(MDP)。形式上，在每一步中，将当前生成的子图视为状态，并将添加一个新的节点或边视为动作。因此，动作概率p(ai|Gi-1)和p(bij|Gi-1,j)可以由上文提到的相应公式来确定。奖励包括中间奖励和最终奖励。中间奖励是对违反价键有效性规则的惩罚。最终奖励包括性质分数和对最终生成分子中的包含过度的空间应变或违规官能团的惩罚。
遵循GCPN的方法，作者使用近端策略优化(PPO)，一个强大的策略优化算法，来优化在上述强化学习环境中的GraphDF模型。对于一个生成的图G，其损失函数为：
其中
是按新策略和旧策略的动作概率的比率。Ai和Aij是估计优势函数，在实验中为未来步骤的累积奖励。
三、实验结果
作者在分子生成的三个任务上评估了GraphDF模型。结果表明，在大多数任务中，GraphDF的性能可以超过三个性能优良的基线，JT-VAE、GCPN和MRNN，以及其他基于流的分子图生成模型，包括GraphNVP、GRF、GraphAF、MoFlow和GraphCNF。
3.1 随机生成
数据集：对于分子图的随机生成，作者在三个分子数据集ZINC250K，QM9，MOSES上训练和评估GraphDF。所有的分子都使用RDKit进行了凯库勒式化。此外，作者还进行了实验，证明了此模型是非常通用的，可以用于分子图以外的其他图。遵从GNF的实验，作者在两个通用的图数据集COMMUNITY-SMALL和EGO-SMALL上测试了GraphDF。
在分子生成方面，作者采用了有效率（Validity），无检查有效率（Validity w/o check），独特性（Uniqueness），新颖性（Novelty）以及重建度（Reconstruction）五个指标进行评估。在通用图数据集上，作者计算了生成的图和数据集中的图之间的度数、聚类系数和轨迹数的最大均值差异（MMD）统计距离。实验结果如下：
表1. 在ZINC数据集上的随机生成性能
表2. 在QM9数据集上的随机生成性能
表3. 在MOSES数据集上的随机生成性能
表4. 在COMMUNITY-SMALL和EGO-SMALL数据集上的随机生成性能
3.2 性质优化
作者采用了惩罚logP指标与定量评估类药性（QED）指标来进行性质优化。作者选取了排名前三的性质分数进行比较。
表5. 通过排名前三的性质分数进行评估的性质优化性能
3.3 约束性质优化
在此任务中，目标为通过修改输入分子图，使得它的惩罚logP分数提高，同时保持优化前后分子的相似性要超过阈值δ。使用的数据集为ZINC250K中惩罚logP分数较低的800个分子。注意，表7和表8中使用的数据集中的800个分子是不同的，为了公平的比较，作者分别在这两个数据集上同相应的模型进行性能比较。实验结果如下：
表7. 在JT-VAE与GCPN中使用的800个分子上进行约束优化的性能比较
表8. 在GraphAF与MoFlow中使用的800个分子上进行约束优化的性能比较
四、总结
作者提出了GraphDF，它是第一个使用离散隐变量的分子生成框架。离散隐变量的使用有利于图结构的更精确的建模和并且具有更强的探索化学空间的能力。作者分析并用实验证明了GraphDF的优点。在未来，作者希望将GraphDF应用于其他类型的图数据，并将其扩展到图翻译问题。
