基于分子圖壓縮的分子生成模型

劉晨陽，劉 勇，惠 麗

(黑龍江大學(xué) 計算機科學(xué)與技術(shù)學(xué)院，哈爾濱 150080)

0 引 言

隨著機器學(xué)習(xí)、深度學(xué)習(xí)等人工智能子領(lǐng)域的發(fā)展以及圖神經(jīng)網(wǎng)絡(luò)的興起[1]，以生成模型為主的深度學(xué)習(xí)技術(shù)在分子設(shè)計中發(fā)揮了極大的作用，以數(shù)據(jù)驅(qū)動的方式探索化學(xué)空間使這一研發(fā)過程變得低成本高效率，同時也提高了研發(fā)過程中的安全性。因此，以深度學(xué)習(xí)技術(shù)解決分子化合物生成問題對于藥物研發(fā)和新分子材料的發(fā)現(xiàn)都有著深遠的意義。

基于深度學(xué)習(xí)技術(shù)分子化合物生成問題通常被定義為輸入一些化合物分子數(shù)據(jù)集，輸出若干個化學(xué)有效的分子，分子互不重復(fù)，且不能存在于訓(xùn)練數(shù)據(jù)集當(dāng)中。以數(shù)據(jù)集訓(xùn)練深度學(xué)習(xí)模型，使被訓(xùn)練的模型擁有足夠的表達能力生成一些滿足輸出要求的分子。分子在計算機中的表達形式通常分為2種，簡化的分子線性輸入規(guī)范(SMILES)字符串和表示為圖結(jié)構(gòu)，且這2種形式之間可互相轉(zhuǎn)換。常用于分子生成任務(wù)的深度學(xué)習(xí)模型有變分自編碼器(VAE)[2]，對抗神經(jīng)網(wǎng)絡(luò)(GAN)[3]，循環(huán)神經(jīng)網(wǎng)絡(luò)(RNN)[4]，以及流式生成模型(Flow)[5-8]。其中RNN主要用于SMILES字符串形式的分子數(shù)據(jù)生成，而對于圖結(jié)構(gòu)形式的分子數(shù)據(jù)，通常以VAE、GAN和Flow建模。

本文采用一種分子壓縮技術(shù)[9]對分子數(shù)據(jù)集進行了壓縮處理，從而實現(xiàn)降低模型或輸入模型維度，受到Zang C 等[10]工作的啟發(fā)提出一種條件流模型(CompMF)，對壓縮后的分子圖建模，將這些分子圖進行編碼，學(xué)習(xí)一個連續(xù)的潛在空間，在潛在空間中采樣隱空間編碼，經(jīng)過流模型的逆過程來生成新穎的有效的分子。該壓縮方法選擇了常見的化學(xué)結(jié)構(gòu)，在降低了模型維度的同時也使模型更容易學(xué)習(xí)到分子構(gòu)成的化學(xué)規(guī)則，該方法在QM9數(shù)據(jù)集上取得了良好的效果。

1 相關(guān)工作

對于分子圖生成工作的研究，主要分為2類:基于簡化的分子線性輸入規(guī)范(SMILES)字符串?dāng)?shù)據(jù)和基于分子圖結(jié)構(gòu)數(shù)據(jù)?；赗NN來設(shè)計生成模型，主要針對SMILES字符串形式的數(shù)據(jù)進行分子生成任務(wù)，例如MRNN[11]和CRNN[12]。其中MRNN采用一種自回歸的方式逐步生成節(jié)點和邊來構(gòu)成最終分子；而CRNN訓(xùn)練了一個從性質(zhì)到分子SMILES字符串的LSTM模型，其生成過程也是由分子的性質(zhì)出發(fā)，逐步生成特定的分子。

而對于分子圖結(jié)構(gòu)數(shù)據(jù)的建模的工作，大都基于VAE[13-20]、GAN和Flow生成模型。其中CGVAE[13]，CCGVAE[14]， NeVAE[15]，GraphVAE[16]，SDVAE[17]，JT-VAE[18]均為基于VAE的分子生成模型，將分子圖數(shù)據(jù)編碼到隱空間，在隱空間中采樣再解碼生成全新的分子圖結(jié)構(gòu)。JT-VAE提出一種新穎的關(guān)于分子骨架的思路，先生成分子的主體骨架，再生成分子的支鏈結(jié)構(gòu)拼裝在骨架上，最后構(gòu)成完整的分子；NeVAE在模型優(yōu)化時，應(yīng)用了貝葉斯優(yōu)化算法對超參數(shù)進行優(yōu)化，取得了良好的效果。

應(yīng)用GAN模型來解決分子生成問題，通常是應(yīng)用一個生成器和一個辨別器來構(gòu)成對抗博弈關(guān)系，反復(fù)迭代直到生成器達到滿意的效果。如Bojchevski A等提出的NetGAN[19]，其生成分子的過程采用了隨機游走的方式；Cao N D等提出的MolGAN[20]在分子生成的新穎性指標(biāo)中表現(xiàn)突出，但在有效性和唯一性上的表現(xiàn)較差一些；You J等提出一種以自回歸的方式生成分子的模型GCPN[21]，它的生成過程是逐個節(jié)點逐條邊依次生成，且這一過程運用了強化學(xué)習(xí)來優(yōu)化，最終生成完整的分子。

對于VAE模型，其優(yōu)化目標(biāo)是誤差的下界，因此基于VAE的模型不夠精確；GAN生成的數(shù)據(jù)，不能通過操縱隱空間向量來控制生成的數(shù)據(jù)，靈活性稍差，F(xiàn)low的提出彌補了VAE和GAN的缺陷。隨著Flow在計算機視覺領(lǐng)域的發(fā)展，如Glow[5]，越來越多的分子生成任務(wù)也采用基于Flow的設(shè)計。Madhawa K等提出的GraphNVP[24]和Honda S等提出的GRF[25]均為基于Flow設(shè)計，可一次性的生成整個分子，但不能很好的保證化學(xué)有效性；Shi C等提出了自回歸流GraphAF[26]，該模型采用自回歸的生成方式，以已生成的節(jié)點和邊為基礎(chǔ)來生成接下來的節(jié)點和邊，并且在邊的添加過程中引入了化學(xué)有效性的檢查與修正，該工作在分子生成任務(wù)中展現(xiàn)了驚人的效果；Zang C等提出了流模型MoFlow[10]可一次性的生成整個分子，而不是以自回歸的方式逐步生成，對分子的整體建模效果更好，且速度較GraphAF更快；Luo Y等提出了離散流模型GraphDF[27]，該模型考慮到圖結(jié)構(gòu)本身就是連續(xù)的，設(shè)計離散流結(jié)構(gòu)減少了反離散化操作，同時也減少了流模型雅可比矩陣的計算開銷。以上這些工作均是直接對分子圖數(shù)據(jù)進行建模，本文采用分子壓縮技術(shù)[9]并提出一個條件流模型(CompMF)，對壓縮后的分子圖進行建模，在有效性、唯一性和新穎性上取得了優(yōu)秀的效果。

2 模型方法

2.1 分子圖壓縮方法

對于一個分子M，其含有的原子個數(shù)為n，則可表示為圖結(jié)構(gòu)G(V,E)，其中V為節(jié)點屬性，表示為n維度的向量，E為邊屬性，表示為[n,n]維度的鄰接矩陣向量。為了后續(xù)的模型訓(xùn)練通常將圖G表示為one-hot向量形式，其節(jié)點矩陣V的維度拓展為[n,ka]，V[i,k]=1表示第i個原子的類型為k；邊矩陣E的維度拓展為[n,n,kb]，E(i,j,c)=1表示第i個原子和第j個原子之間的化學(xué)鍵類型為c。其中ka為分子數(shù)據(jù)集中的原子種類數(shù)；kb為分子數(shù)據(jù)集中的化學(xué)鍵種類數(shù)(一般僅包含單間、雙鍵、三鍵)。采用Kwon Y等提出的分子壓縮方法，選取6種分子子結(jié)構(gòu)(圖1)，分別為-C-C-結(jié)構(gòu)、-C-結(jié)構(gòu)、-C=C-結(jié)構(gòu)、=C-C=結(jié)構(gòu)、-N-結(jié)構(gòu)和-O-結(jié)構(gòu)。該6種子結(jié)構(gòu)在分子圖數(shù)據(jù)中最為常見，可保證更多的分子盡可能大的壓縮，若分子中包含這些子結(jié)構(gòu)，則用新定義的邊代替子結(jié)構(gòu)，邊的種類增多，但分子圖的節(jié)點數(shù)目減少，分子圖被有效壓縮。該壓縮方法的壓縮效果見圖2。經(jīng)過壓縮減少了2個原子的分子多達46 650個，所占比重最多，減少了3個原子的分子占比第二多，而分子最多減少了6個原子，且有5 765個分子沒有被壓縮，整體的壓縮率為95.694%，大部分的壓縮集中在縮減了2個和3個原子維度上。

圖1 定義的6個子結(jié)構(gòu)Fig.1 Definition of six substructures

圖2 壓縮后減少的維度分布Fig.2 Dimension distribution reduce after compressing

圖3 分子圖壓縮展示Fig.3 Compression of molecular graph

6種子結(jié)構(gòu)由SMART字符串定義，以rdkit庫轉(zhuǎn)化為特定的分子結(jié)構(gòu)對象，在分子圖中匹配查找，按照上述規(guī)律壓縮后的分子依舊為圖結(jié)構(gòu)形式G(V,E)，其kb的值變大，但n的值隨之變小。壓縮過程見圖3，左圖為甲苯的原始分子結(jié)構(gòu)，右圖為經(jīng)過壓縮后的甲苯分子結(jié)構(gòu)，右圖中的紅邊對應(yīng)子結(jié)構(gòu)3，右圖中的藍邊對應(yīng)子結(jié)構(gòu)4?？芍庇^的看到甲苯原有7個重原子，經(jīng)過壓縮后，壓縮為3個重原子，降低了分子中的重原子數(shù)。若在壓縮過程中發(fā)生結(jié)構(gòu)重疊的現(xiàn)象，則以重原子序數(shù)之和大的子結(jié)構(gòu)優(yōu)先。如在壓縮的過程中，發(fā)生子結(jié)構(gòu)1和子結(jié)構(gòu)5重合，由于子結(jié)構(gòu)1的重原子序數(shù)之和為12，子結(jié)構(gòu)5的重原子序數(shù)之和為7，此時選擇以子結(jié)構(gòu)1來壓縮。

由于模型以經(jīng)過壓縮后的分子圖來訓(xùn)練，則在生成分子的環(huán)節(jié)，模型生成的分子圖也是經(jīng)過壓縮后的，所以模型輸出的生成分子要按照上述的壓縮方法來還原，將6種子結(jié)構(gòu)的對應(yīng)邊替換為原始的6種子結(jié)構(gòu)，形成最終的分子結(jié)構(gòu)圖。

2.2 條件流模型

流模型是一種可逆的神經(jīng)網(wǎng)絡(luò)可實現(xiàn)一種數(shù)據(jù)分布到另一種數(shù)據(jù)分布的可逆映射，根據(jù)變量替換式(1)，實現(xiàn)真實數(shù)據(jù)分布PX(X)到已知的簡單的數(shù)據(jù)PZ(Z)分布的可逆映射。

(1)

由于一次變換模型的映射能力不夠精確，實際應(yīng)用中設(shè)計一系列的可逆變換函數(shù)，即

(2)

其中,K為流模塊的個數(shù)，其對數(shù)似然形式:

(3)

對于分子圖數(shù)據(jù)G(V,E)，CompMF模型分為2個流模型：邊流fB(B)和給定邊后的節(jié)點條件流fA|B(A|B)。邊流fB(B)用來建立邊的特征數(shù)據(jù)E和隱空間分布的可逆映射，節(jié)點條件流fA|B(A|B)根據(jù)給出的邊數(shù)據(jù)建立節(jié)點特征V與隱空間分布的可逆映射。邊流與節(jié)點條件流的放射耦合過程均采用Glow[5]中的放射耦合方式，將輸入數(shù)據(jù)分為兩個維度相等的部分Zh和Zl或Xh和Xl，經(jīng)過簡單運算實現(xiàn)可逆變換。

Zh=Xh

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

(4)

Xh=Zh

(5)

對于分子圖結(jié)構(gòu)數(shù)據(jù)，Sθ(·)和Tθ(·)為圖卷積網(wǎng)絡(luò)。CompMF模型框架見圖4，其中node為節(jié)點特征矩陣;bond為邊特征矩陣;compress為分子的壓縮過程；extend為分子的解壓還原過程。

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

邊流f(B)將bond數(shù)據(jù)可逆映射到ZB，即ZB=fB(B)，如式(3)可得損失函數(shù)為

(6)

節(jié)點流f(A|B)用于建模給定邊數(shù)據(jù)的節(jié)點數(shù)據(jù)信息，將node數(shù)據(jù)以給定bond信息的條件分布映射到ZA|B，即ZA|B|B=fA|B(A|B)，損失函數(shù)為

(7)

2.3 化合價修正

由于流模型輸出的僅為分子的節(jié)點特征矩陣V和邊屬性矩陣E，兩個張量重構(gòu)成一個分子可能會出現(xiàn)化學(xué)無效的情況，要對生成的分子做化學(xué)有效性檢查與修正。本文采用MoFlow[10]中的化學(xué)有效性修正方法，其修正規(guī)則：

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

(8)

c為邊的種類，取值為{1,2,3};E(i,j,kb)為邊屬性one-hot張量;Ch為形式電荷。

首先對分子圖進行壓縮，然后對壓縮后的圖數(shù)據(jù)建模，生成的分子經(jīng)過反壓縮過程的展開得到其對應(yīng)的正確形式，最后進行化學(xué)有效性的檢查與修正，得到新穎的有效的分子。

3 實 驗

采用QM9數(shù)據(jù)集進行實驗，該數(shù)據(jù)集擁有約13萬個分子，最大重原子數(shù)為9，包含C,N,O,F重原子。

3.1 實驗過程

根據(jù)圖1中的子結(jié)構(gòu)對分子數(shù)據(jù)集中的每個分子進行壓縮處理，將壓縮后的分子表示為one-hot向量形式的張量，節(jié)點表示為[n,ka]維度的張量，邊表示為[n,n,kb]維度的張量，以此類型的數(shù)據(jù)對條件流模型進行訓(xùn)練。

在pytorch2.7環(huán)境下，采用Adam自適應(yīng)優(yōu)化器優(yōu)化，設(shè)置初始學(xué)習(xí)率為0.005，batchsize大小為256，以QM9數(shù)據(jù)集在RTX3090(cuda11.4)上訓(xùn)練了200個周期，大約用時4 h。

在測試階段，從模型的隱空間隨機采樣10 000個隱向量z，將隱向量z送入流模型中，輸出了10 000個分子，分子為經(jīng)過壓縮后的形式，采用壓縮方法的逆替換過程，將其擴展為最終的分子。在有效性、唯一性、新穎性3個指標(biāo)上評估了這些分子，并展示了其分子結(jié)構(gòu)，見圖5。

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

3.2 實驗結(jié)果

對比了近兩年的分子圖生成研究成果，結(jié)果見表1，其中GraphNVP，GRF，GraphAF，MoFlow，GraphDF均為基于流式生成模型的方法，直接對分子圖數(shù)據(jù)進行建模，均未采用分子圖壓縮技術(shù)。有效性描述生成的分子當(dāng)中化學(xué)正確的分子的比率；唯一性描述生成的分子之間相互不重復(fù)的比率；新穎性描述生成的分子與訓(xùn)練數(shù)據(jù)集的不重復(fù)程度。表達式為

表1 QM9數(shù)據(jù)集實驗結(jié)果對比

(9)

(10)

(11)

其中，train為訓(xùn)練數(shù)據(jù)集；gens表示生成的分子集合；num(·)為計數(shù)函數(shù)；valid(·)返回分子集合中化學(xué)有效的分子集合；set(·)為去重函數(shù)。

由表1可見，在化學(xué)有效性(Validity)指標(biāo)中，模型CompMF和GraphAF[26]，MoFlow[10]以及GraphDF[27]都已經(jīng)達到了100%，而CompMF在唯一性(Uniqueness)和新穎性(Novelty)兩項指標(biāo)中也要比GraphAF[26]，MoFlow[10]和GraphDF[27]表現(xiàn)得更好。

4 結(jié) 論

本文采用分子圖壓縮技術(shù)，并提出一個條件流模型CompMF對經(jīng)過壓縮后的分子圖進行建模，在分子生成工作中的表現(xiàn)優(yōu)異，同時采用分子壓縮技術(shù)后的流模型維度更低，減輕了雅克比行列式的計算消耗。表示邊的鄰接矩陣普遍要比表示節(jié)點特征的矩陣所包含的數(shù)據(jù)量要更多，本文采用的壓縮方法是將6種常見的子結(jié)構(gòu)定義為不同類型的邊，雖然大幅降低了節(jié)點的數(shù)量，但也增加了邊的數(shù)量，數(shù)據(jù)量的壓縮程度并不是很高，可嘗試將子結(jié)構(gòu)定義為不同類型的節(jié)點；此外，在化學(xué)分子結(jié)構(gòu)中有一些常見的官能團，如羧基，醛基，羰基，醚鍵等，是分子中常見的子結(jié)構(gòu)且更能代表一些化學(xué)性質(zhì)，嘗試基于這些官能團子結(jié)構(gòu)來壓縮分子圖將會更具有意義。