基于分子图压缩的分子生成模型

刘晨阳，刘 勇，惠 丽

(黑龙江大学 计算机科学与技术学院，哈尔滨 150080)

0 引 言

随着机器学习、深度学习等人工智能子领域的发展以及图神经网络的兴起[1]，以生成模型为主的深度学习技术在分子设计中发挥了极大的作用，以数据驱动的方式探索化学空间使这一研发过程变得低成本高效率，同时也提高了研发过程中的安全性。因此，以深度学习技术解决分子化合物生成问题对于药物研发和新分子材料的发现都有着深远的意义。

基于深度学习技术分子化合物生成问题通常被定义为输入一些化合物分子数据集，输出若干个化学有效的分子，分子互不重复，且不能存在于训练数据集当中。以数据集训练深度学习模型，使被训练的模型拥有足够的表达能力生成一些满足输出要求的分子。分子在计算机中的表达形式通常分为2种，简化的分子线性输入规范(SMILES)字符串和表示为图结构，且这2种形式之间可互相转换。常用于分子生成任务的深度学习模型有变分自编码器(VAE)[2]，对抗神经网络(GAN)[3]，循环神经网络(RNN)[4]，以及流式生成模型(Flow)[5-8]。其中RNN主要用于SMILES字符串形式的分子数据生成，而对于图结构形式的分子数据，通常以VAE、GAN和Flow建模。

本文采用一种分子压缩技术[9]对分子数据集进行了压缩处理，从而实现降低模型或输入模型维度，受到Zang C 等[10]工作的启发提出一种条件流模型(CompMF)，对压缩后的分子图建模，将这些分子图进行编码，学习一个连续的潜在空间，在潜在空间中采样隐空间编码，经过流模型的逆过程来生成新颖的有效的分子。该压缩方法选择了常见的化学结构，在降低了模型维度的同时也使模型更容易学习到分子构成的化学规则，该方法在QM9数据集上取得了良好的效果。

1 相关工作

对于分子图生成工作的研究，主要分为2类:基于简化的分子线性输入规范(SMILES)字符串数据和基于分子图结构数据。基于RNN来设计生成模型，主要针对SMILES字符串形式的数据进行分子生成任务，例如MRNN[11]和CRNN[12]。其中MRNN采用一种自回归的方式逐步生成节点和边来构成最终分子；而CRNN训练了一个从性质到分子SMILES字符串的LSTM模型，其生成过程也是由分子的性质出发，逐步生成特定的分子。

而对于分子图结构数据的建模的工作，大都基于VAE[13-20]、GAN和Flow生成模型。其中CGVAE[13]，CCGVAE[14]， NeVAE[15]，GraphVAE[16]，SDVAE[17]，JT-VAE[18]均为基于VAE的分子生成模型，将分子图数据编码到隐空间，在隐空间中采样再解码生成全新的分子图结构。JT-VAE提出一种新颖的关于分子骨架的思路，先生成分子的主体骨架，再生成分子的支链结构拼装在骨架上，最后构成完整的分子；NeVAE在模型优化时，应用了贝叶斯优化算法对超参数进行优化，取得了良好的效果。

应用GAN模型来解决分子生成问题，通常是应用一个生成器和一个辨别器来构成对抗博弈关系，反复迭代直到生成器达到满意的效果。如Bojchevski A等提出的NetGAN[19]，其生成分子的过程采用了随机游走的方式；Cao N D等提出的MolGAN[20]在分子生成的新颖性指标中表现突出，但在有效性和唯一性上的表现较差一些；You J等提出一种以自回归的方式生成分子的模型GCPN[21]，它的生成过程是逐个节点逐条边依次生成，且这一过程运用了强化学习来优化，最终生成完整的分子。

对于VAE模型，其优化目标是误差的下界，因此基于VAE的模型不够精确；GAN生成的数据，不能通过操纵隐空间向量来控制生成的数据，灵活性稍差，Flow的提出弥补了VAE和GAN的缺陷。随着Flow在计算机视觉领域的发展，如Glow[5]，越来越多的分子生成任务也采用基于Flow的设计。Madhawa K等提出的GraphNVP[24]和Honda S等提出的GRF[25]均为基于Flow设计，可一次性的生成整个分子，但不能很好的保证化学有效性；Shi C等提出了自回归流GraphAF[26]，该模型采用自回归的生成方式，以已生成的节点和边为基础来生成接下来的节点和边，并且在边的添加过程中引入了化学有效性的检查与修正，该工作在分子生成任务中展现了惊人的效果；Zang C等提出了流模型MoFlow[10]可一次性的生成整个分子，而不是以自回归的方式逐步生成，对分子的整体建模效果更好，且速度较GraphAF更快；Luo Y等提出了离散流模型GraphDF[27]，该模型考虑到图结构本身就是连续的，设计离散流结构减少了反离散化操作，同时也减少了流模型雅可比矩阵的计算开销。以上这些工作均是直接对分子图数据进行建模，本文采用分子压缩技术[9]并提出一个条件流模型(CompMF)，对压缩后的分子图进行建模，在有效性、唯一性和新颖性上取得了优秀的效果。

2 模型方法

2.1 分子图压缩方法

对于一个分子M，其含有的原子个数为n，则可表示为图结构G(V,E)，其中V为节点属性，表示为n维度的向量，E为边属性，表示为[n,n]维度的邻接矩阵向量。为了后续的模型训练通常将图G表示为one-hot向量形式，其节点矩阵V的维度拓展为[n,ka]，V[i,k]=1表示第i个原子的类型为k；边矩阵E的维度拓展为[n,n,kb]，E(i,j,c)=1表示第i个原子和第j个原子之间的化学键类型为c。其中ka为分子数据集中的原子种类数；kb为分子数据集中的化学键种类数(一般仅包含单间、双键、三键)。采用Kwon Y等提出的分子压缩方法，选取6种分子子结构(图1)，分别为-C-C-结构、-C-结构、-C=C-结构、=C-C=结构、-N-结构和-O-结构。该6种子结构在分子图数据中最为常见，可保证更多的分子尽可能大的压缩，若分子中包含这些子结构，则用新定义的边代替子结构，边的种类增多，但分子图的节点数目减少，分子图被有效压缩。该压缩方法的压缩效果见图2。经过压缩减少了2个原子的分子多达46 650个，所占比重最多，减少了3个原子的分子占比第二多，而分子最多减少了6个原子，且有5 765个分子没有被压缩，整体的压缩率为95.694%，大部分的压缩集中在缩减了2个和3个原子维度上。

图1 定义的6个子结构Fig.1 Definition of six substructures

图2 压缩后减少的维度分布Fig.2 Dimension distribution reduce after compressing

图3 分子图压缩展示Fig.3 Compression of molecular graph

6种子结构由SMART字符串定义，以rdkit库转化为特定的分子结构对象，在分子图中匹配查找，按照上述规律压缩后的分子依旧为图结构形式G(V,E)，其kb的值变大，但n的值随之变小。压缩过程见图3，左图为甲苯的原始分子结构，右图为经过压缩后的甲苯分子结构，右图中的红边对应子结构3，右图中的蓝边对应子结构4。可直观的看到甲苯原有7个重原子，经过压缩后，压缩为3个重原子，降低了分子中的重原子数。若在压缩过程中发生结构重叠的现象，则以重原子序数之和大的子结构优先。如在压缩的过程中，发生子结构1和子结构5重合，由于子结构1的重原子序数之和为12，子结构5的重原子序数之和为7，此时选择以子结构1来压缩。

由于模型以经过压缩后的分子图来训练，则在生成分子的环节，模型生成的分子图也是经过压缩后的，所以模型输出的生成分子要按照上述的压缩方法来还原，将6种子结构的对应边替换为原始的6种子结构，形成最终的分子结构图。

2.2 条件流模型

流模型是一种可逆的神经网络可实现一种数据分布到另一种数据分布的可逆映射，根据变量替换式(1)，实现真实数据分布PX(X)到已知的简单的数据PZ(Z)分布的可逆映射。

(1)

由于一次变换模型的映射能力不够精确，实际应用中设计一系列的可逆变换函数，即

(2)

其中,K为流模块的个数，其对数似然形式:

(3)

对于分子图数据G(V,E)，CompMF模型分为2个流模型：边流fB(B)和给定边后的节点条件流fA|B(A|B)。边流fB(B)用来建立边的特征数据E和隐空间分布的可逆映射，节点条件流fA|B(A|B)根据给出的边数据建立节点特征V与隐空间分布的可逆映射。边流与节点条件流的放射耦合过程均采用Glow[5]中的放射耦合方式，将输入数据分为两个维度相等的部分Zh和Zl或Xh和Xl，经过简单运算实现可逆变换。

Zh=Xh

Zl=XleSθ(Xh)+Tθ(Xh)

(4)

Xh=Zh

(5)

对于分子图结构数据，Sθ(·)和Tθ(·)为图卷积网络。CompMF模型框架见图4，其中node为节点特征矩阵;bond为边特征矩阵;compress为分子的压缩过程；extend为分子的解压还原过程。

图4 CompMF模型框架Fig.4 Model of CompMF

边流f(B)将bond数据可逆映射到ZB，即ZB=fB(B)，如式(3)可得损失函数为

(6)

节点流f(A|B)用于建模给定边数据的节点数据信息，将node数据以给定bond信息的条件分布映射到ZA|B，即ZA|B|B=fA|B(A|B)，损失函数为

(7)

2.3 化合价修正

由于流模型输出的仅为分子的节点特征矩阵V和边属性矩阵E，两个张量重构成一个分子可能会出现化学无效的情况，要对生成的分子做化学有效性检查与修正。本文采用MoFlow[10]中的化学有效性修正方法，其修正规则：

∑cE(i,j,kb)≤valency(atomsi)+Ch

(8)

c为边的种类，取值为{1,2,3};E(i,j,kb)为边属性one-hot张量;Ch为形式电荷。

首先对分子图进行压缩，然后对压缩后的图数据建模，生成的分子经过反压缩过程的展开得到其对应的正确形式，最后进行化学有效性的检查与修正，得到新颖的有效的分子。

3 实 验

采用QM9数据集进行实验，该数据集拥有约13万个分子，最大重原子数为9，包含C,N,O,F重原子。

3.1 实验过程

根据图1中的子结构对分子数据集中的每个分子进行压缩处理，将压缩后的分子表示为one-hot向量形式的张量，节点表示为[n,ka]维度的张量，边表示为[n,n,kb]维度的张量，以此类型的数据对条件流模型进行训练。

在pytorch2.7环境下，采用Adam自适应优化器优化，设置初始学习率为0.005，batchsize大小为256，以QM9数据集在RTX3090(cuda11.4)上训练了200个周期，大约用时4 h。

在测试阶段，从模型的隐空间随机采样10 000个隐向量z，将隐向量z送入流模型中，输出了10 000个分子，分子为经过压缩后的形式，采用压缩方法的逆替换过程，将其扩展为最终的分子。在有效性、唯一性、新颖性3个指标上评估了这些分子，并展示了其分子结构，见图5。

图5 生成的分子图Fig.5 Molecular graphs of generation

3.2 实验结果

对比了近两年的分子图生成研究成果，结果见表1，其中GraphNVP，GRF，GraphAF，MoFlow，GraphDF均为基于流式生成模型的方法，直接对分子图数据进行建模，均未采用分子图压缩技术。有效性描述生成的分子当中化学正确的分子的比率；唯一性描述生成的分子之间相互不重复的比率；新颖性描述生成的分子与训练数据集的不重复程度。表达式为

表1 QM9数据集实验结果对比

(9)

(10)

(11)

其中，train为训练数据集；gens表示生成的分子集合；num(·)为计数函数；valid(·)返回分子集合中化学有效的分子集合；set(·)为去重函数。

由表1可见，在化学有效性(Validity)指标中，模型CompMF和GraphAF[26]，MoFlow[10]以及GraphDF[27]都已经达到了100%，而CompMF在唯一性(Uniqueness)和新颖性(Novelty)两项指标中也要比GraphAF[26]，MoFlow[10]和GraphDF[27]表现得更好。

4 结 论

本文采用分子图压缩技术，并提出一个条件流模型CompMF对经过压缩后的分子图进行建模，在分子生成工作中的表现优异，同时采用分子压缩技术后的流模型维度更低，减轻了雅克比行列式的计算消耗。表示边的邻接矩阵普遍要比表示节点特征的矩阵所包含的数据量要更多，本文采用的压缩方法是将6种常见的子结构定义为不同类型的边，虽然大幅降低了节点的数量，但也增加了边的数量，数据量的压缩程度并不是很高，可尝试将子结构定义为不同类型的节点；此外，在化学分子结构中有一些常见的官能团，如羧基，醛基，羰基，醚键等，是分子中常见的子结构且更能代表一些化学性质，尝试基于这些官能团子结构来压缩分子图将会更具有意义。