改進的語言模型和循環神經網絡優化新藥設計

陳 相

（武漢科技大學計算機科學與技術學院，武漢 430070）

0 引言

作為一個全球性的學術和商業過程，藥物發現是藥物設計過程的第一步。然而，大規模探索新的化學空間在支持藥物發現方面一直是一項艱巨的任務和挑戰。在過去的幾十年里，機器學習越來越多地被用于制造有意義的分子［1-2］。通過從零開始生成分子，從頭開始藥物設計可以更好地表達整個化學空間，包括深度強化學習（deep reinforcement learning，DRL）在內的機器學習已成功用于從頭開始藥物設計方法的開發［3-4］，從而達到小化學空間搜索范圍的目的，使研究人員能夠專注于感興趣的領域。

1 優化模型的整體框架

1.1 循環神經網絡

循環神經網絡（recurrent neural network，RNN）的結構可以很好地利用序列與時序之間的關系以及數據中的語義信息，是處理序列特性數據的有效方法。因此，在分子信息學領域，RNN成功地解決了機器學習任務，被廣泛應用于計算分子生成。在最近的一系列實驗中［5-7］，基于簡化分子輸入行輸入系統（simplified molecular input line entry systems,SMILES）的具有長短期記憶（long short-term memory，LSTM）細胞的RNN在捕獲蛋白質分子的結構和功能以及產生預測類似性質的分子方面取得了很好的結果。使用SMILES 將大量已知的活性化合物編碼為SMILES 字符串序列。LSTM 模型基于自然語言處理，使用SMILES 字符序列（“標記”）作為輸入，根據給定的一系列先前字符和概率估計學習預測下一個SMILES 字符。實驗表明，該方法在生成高質量的新分子方面是非常有效的。

1.2 自然語言處理

在使用自然語言處理方法的過程中，主要工作是表征學習［8］，自然語言處理中最常用的符號化方法之一是詞向量［9］。最早的詞向量采用獨熱編碼，將詞轉化為一個很長的稀疏向量。向量的維度等于詞匯表的大小，向量中只有一個維度為1，其他維度均為0。這種編碼方法簡潔，但一個重要的缺陷是忽略了文本中單詞的順序，編碼中所有的單詞向量都是相互正交的，沒有體現單詞之間的相似關系。受分布式假設（即相似環境中的單詞通常具有相似的含義）的啟發，Word2Vec 模型［10］采用分布式表示有效地解決了獨熱編碼問題，其思路是通過訓練將原本由獨熱編碼的每個單詞映射成更短的單詞向量，這樣就可以輕松地分析單詞之間的關系，使NLP 問題更容易解決。圖1 展示了本實驗模型的簡單流程，Word2Vec 語言模型作為預訓練模型，用于訓練由SMILES 編碼的文本詞匯表的表示向量。并將它們整合到循環神經網絡中，以訓練和產生新的分子。

圖1 實驗模型流程

知識蒸餾的過程是將一個大型、表現良好的模型（教師模型）的知識轉移給一個較小的模型（學生模型），可以看作是模型壓縮的一種形式，其中學生模型被訓練成模仿教師模型的行為［11-13］。通過利用教師模型預測中編碼的豐富知識，學生模型可以從其優越的性能中受益，同時保持緊湊的大小，允許學生模型更好地泛化。

在本研究中，為了提高蛋白質表達能力，將SMILES 編碼的氨基酸序列劃分為單詞（k-gram），輸入到Word2Vec 模型中，得到單詞的分布式表示向量。實驗結果表明，通過對Word2Vec 模型和LSTM 模型進行優化和調整，改進后的模型能夠更好地從蛋白質特征空間中學習特征，與傳統模型相比取得了優異的性能，生成的有效分子百分比比傳統模型提高了15 個百分點，分子質量相比于傳統模型大大提高，并且在微調訓練模型時，我們通過知識蒸餾利用少量數據實現大模型到小模型的知識遷移，結果表明，小模型可以很好地學習到大模型的能力，達到較好的效果。

2 方法

2.1 數據預處理

我們從開源的CHEMBL 數據集中組裝一個分子的訓練數據集，并用SMILES 編碼表示分子，通過對數據集進行預處理去除重復的化學信息，并過濾我們采用的化學空間之外的分子，最終得到近500000 個分子的數據集。為便于RNN 模型訓練，我們用‘G’（GO）和‘ ’（換行）作為每個SMILES 分子開始和結束的標記符，所有分子長度在35～75 個字符之間，最終產生一個共53個獨特字符的詞匯表W。

2.2 預訓練模型

Word2Vec 是一種廣泛應用于自然語言處理的語言模型，一種以大型本文數據作為輸入，從大量文本中以無監督方式學習語義知識的模型。對詞匯表中的任意一個單詞w，可以生成一個固定長度的實值向量V(w)，所有的詞向量可看作空間中的一個點，它們之間的語義相似度可以通過點之間的“距離”來判斷。我們希望模型訓練的結果能夠讓相似語義的單詞具有相似的表征向量，近年來，這種表示能夠在藥物設計等問題上帶來很好的結果，例如藥物-靶標相互作用預測特征表示SPVec［14］、從上下文感知中學習分子表示Mol2Context-vec［15］。訓練Word2Vec 的常用方法有連續詞袋（Continuous Bag-of-Words，CBOW）模型和跳躍圖模型（Skip-Gram）［16］。圖2 使用谷歌的開源嵌入投影儀工具將詞匯表W 的詞向量映射到三維空間，并輸出與字符距離最近的十個字符。

圖2 詞向量的空間分布

我們將經過預處理的分子數據集作為無監督訓練語料，由于標準的SMILES 具有唯一性，每個SMILES 編碼能夠對應唯一一個化學結構，如圖3 所示。因此將每一條SMILES 字符串作為“句子”，將其劃分成一系列的單詞（k-grams）。結合RNN 模型訓練數據的特征，經過實驗對比，最后可以得到當k為1 時，可以最大程度地保證氨基酸分子的化學特征。

圖3 分子示例及其SMILES表示

我們使用負采樣優化方法的Skip-Grams模型來訓練Word2Vec 模型，它的算法在于給定中心詞的情況下，預測它的上下文單詞，相比于CBOW 模型單詞之間有更多的訓練機會［17］。經典的Skip-Gram 模型由輸入層、投影層、輸出層三部分組成。如圖4 所示，輸入層以(w,Postive(w),Negtive(w))作為一個樣本輸入，其中w是選定的中心詞，Postive（w）由中心詞w前后C個詞組成，C是模型上下文窗口大小。Negtive（w）是w的K個負例，它是以單詞出現的頻率為權重通過負采樣的方式隨機生成的，正例與負例統稱為背景詞。我們將Skip-Gram 模型的超參數設置為上下文窗口C=5，負采樣比例a= 10 (K=a*C)，嵌入矩陣維度D= 100，樣本每次以B= 128 大小的批次輸入網絡。將中心詞w的獨熱向量乘以權矩陣W1，得到中心詞w向量表達式U。

圖4 Skip-Gram 模型

同樣，將中心詞w選取的背景詞的獨熱向量乘以權矩陣W2，得到背景詞樣本V的向量表達式。

接下來，輸出Y表示從中心詞w預測每個背景詞的概率，其值范圍為（-1,1）。

在這個模型中，隱層不使用任何激活函數，但在輸出層我們改變傳統的sigmoid 函數，使用tanh 激活函數將輸出向量的所有值映射到（-1, 1），同時定義一個向量label 作為標簽，其中正例為1，負例為-1。則有：

針對單個輸出，p(wj|w)表示模型輸出為wj的概率，模型的損失函數對于所有的正例和負例的輸出值，輸出正確的概率P為：

這里是yj對應向量label 中的標簽，最大化輸出單詞概率P，即最小化損失函數L：

可以使用隨機梯度下降技術來最小化損失函數。經過100 次迭代訓練后，我們將訓練得到的詞嵌入矩陣放入LSTM 模型中作為輸入分子的初始化權重。

2.3 LSTM模型

循環神經網絡（RNN）可以對輸入的序列數據建模，給定時間步長為t的序列X=x1x2x3…xt和對應的輸出序列Y=y1y2y3…yt，模型訓練過程由一個字符預測下一個字符，因此定義yi=xi+1。任一時間步長的xi，RNN 通過一系列的門返回輸出概率，并將隱藏狀態hi傳遞給下一個單元。由此可見由前i個時間步長的序列共同決定。

LSTM 是為了解決因序列較長而導致訓練過程中產生的梯度消失和梯度爆炸問題而提出的一種特殊RNN 模型，引入門控的選擇性機制，如圖5 所示，每個LSTM 單元都有稱為遺忘門、輸入門、輸出門的神經網絡層，能夠有選擇性地遺忘和保留訓練中的信息。

圖5 LSTM的內部機制和相關參數

如圖6 所示，本文模型由三個LSTM 層組成，每個層的隱藏狀態向量為512，并使用dropout 進行正則化。LSTM 層前面是一層嵌入層，其權重是在Word2Vec 模型中訓練得到的詞嵌入矩陣，后面是一層密集輸出層和一個Softmax激活函數的神經元。模型以75 為時間步長，128批次輸入，序列經過嵌入層獲取詞向量后進入LSTM 單元。前向傳播過程中，我們將一個序列最后時間步長的單元格狀態和隱藏狀態直接作為下一序列的初始化單元格和隱藏狀態，每128 批次的序列只初始化一次LSTM的隱藏狀態，后面的序列可以使用前面序列訓練過程中保存下來的信息。應用密集層以產生輸出logits，然后通過Softmax 層將其轉換為概率。利用交叉熵損失函數和Adam優化算法進行性能優化。

圖6 LSTM模型訓練流程

2.4 分子評估與模型微調

藥物設計過程中我們需要對生成的分子特性［18］進行評估，確定滿足我們特性需求的藥物分子，以便進行下一步的研究。本實驗在分子評估過程中借鑒Yasonik［5］進行優化后的三法則（RO3），即辛醇-水分配系數logP≤3、分子量≤480 g/mol、≤3 個氫鍵供體、≤3 個氫鍵受體、≤3 個可旋轉鍵，使用化學信息學庫RDKit實現分子評估，并使用Fonseca 和Fleming 的非支配排序算法來比較模型根據RO3 的標準生成的分子。對LSTM 網絡進行訓練生成有效的SMILES 字符串后，我們需要再選定一個較小的數據集進一步訓練模型來對模型進行微調，使模型調整為能夠生成與有效分子集具有更高相似性的SMILES 字符串，這一過程實際上使用遷移學習方法，能夠用較少的數據集優選分子特定屬性。實驗中，我們將基于優化后的RO3 的五個約束的非支配排序算法選擇在生成的有效分子中最好的一半作為新的數據集對模型進行微調，選定的分子被輸入LSTM 模型，再次訓練并生成新的有效分子。每一次微調后都會對新生成的有效分子進行評估并再選取最好的一半做下一次微調。最后，我們用大模型作為教師模型，指導只有一層LSTM 的學生模型訓練生成有小分子，可以看到在少量數據訓練后，學生模型在生成獨特且新穎的有效分子方面有很好的效果。

3 實驗結果分析

將50 萬個分子的數據集輸入模型進行訓練，從訓練后的LSTM 網絡中生成100 萬個SMILES 字符，在迭代微調之前，模型共產生了21294個分子。我們將最終實驗結果與Yasonik［5］中的神經網絡模型訓練生成的分子比較。結果表明，相比于傳統的循環神經網絡模式，使用優化后的Word2Vec 模型與LSTM 模型結合可以有效提高藥物分子生成的能力。如圖7 所示。我們用RDKIT 評估分子的有效性和相關特征，能夠從三個方面證明其高效性：①有效分子百分比提高，從傳統模型中采樣的100 萬個SMILES 字符，生成的19722 個分子中，77%是有效的，本實驗改進模型后，有效分子百分比提高至92%；②獨特且新穎的有效分子數提高，傳統模型產生的有效分子過濾無效和重復的分子后留下了9415 個獨特、新穎和有效的分子，本實驗模型最終獨特、新穎和有效的分子達到了12850 個；③訓練效率提高，由于詞向量能將序列編碼為更緊湊的向量，降低內存的占用，模型訓練的時間大大減少，本實驗模型訓練的總時間不到傳統模型的三分之一。

圖7 實驗模型相比于傳統模型的高效性

為了進一步比較實驗模型與傳統模型之間的差異，我們對生成的藥物分子做更詳細的評估，為此，我們從50 萬個分子的數據集中截取5萬個分子，同時輸入實驗的優化模型和傳統模型中訓練，并對模型進行迭代微調，每次訓練結束后都從模型中采樣100 萬個SMILES 字符，對生成的分子進行評估。

我們對模型進行了迭代微調，從圖8 可以發現模型在第三次迭代時，改進后的模型能夠以85%左右的穩定效率生成分子，新穎獨特的有效分子數量從5600 個增加到11700 個，與傳統模型相比，效果明顯提高。

圖8 模型迭代結果

根據RO3 評價標準，從兩種模型生成的有效分子中選出最優的5000 個分子，再次通過PCA進行分子性質可視化，如圖9和圖10所示。

圖9 改良模型的分子更多滿足RO3的評估標準

圖10 傳統模型和改良模型的屬性分布

隨后，我們從生成的有效分子中精選出8000 條獨特且新穎的有效分子，用改進的模型作為教師模型，訓練具有一層LSTM 的學生模型生成有效分子，學生模型可以接近三層LSTM 模型的能力，并在生成獨立且新穎的有效分子上超越了三層LSTM 模型，達到模型壓縮的效果。在實驗中，我們將傳統模型中的sigmoid 函數替換為tanh 函數作為輸出層的激活函數，并對參數進行了優化，調整了損失函數和模型反向傳播。從圖11 的結果可以看出，這種改進對模型的優化起到了很大的作用。

圖11 sigmoid函數的傳統模型和tanh函數改進模型對比

基于RO3 評估分子，我們也用PCA 可視化兩個模型生成的分子中最優的5000 個分子的屬性，如圖12 和圖13，并分別可視化5 個屬性?？梢钥闯鍪褂胻anh 函數的模型生成的分子有更高的質量。

圖12 tanh激活函數模型的分子更多滿足RO3評估標準

圖13 sigmoid和tanh激活函數模型分子的屬性分布（續）

圖13 sigmoid和tanh激活函數模型分子的屬性分布

4 結語

本文將詞向量模型與循環神經網絡相結合，將蛋白質分子中隱含的特征緊密聯系在一起，Word2Vec 嵌入矩陣可視化很好地說明了具有相似特性的SMILES 字符向量緊密地位于向量空間中。通過對Word2Vec 和LSTM 模型反復優化，包括模型權重、超參數、輸入序列、損失函數等結構進行多次調整和比較，最終分析一組最佳的組合能夠使模型在藥物生成實驗中達到最高效的水平。

最初，LSTM 模型使用獨熱編碼方法處理輸入序列，把SMILES 字符轉化成一個很長的稀疏向量，與獨熱編碼方法相比，Word2Vec 模型通過訓練語言模型將輸入序列編碼為更緊湊的數字向量，既有更佳的表達能力，又能減小內存的使用。實驗最終證明這一事實，將Word2Vec作為預訓練模型為LSTM 模型的訓練數據集做第一步訓練，使得LSTM 模型在訓練時間和生成分子有效性方面都有大幅度的優化，并通過知識蒸餾可以對模型進行壓縮，而不損失模型的效果。

總而言之，我們認為類似Word2Vec 的自然語言處理方法和知識蒸餾技術在新藥物設計上有相當廣泛的運用空間，在降低與藥物開發相關的成本和時間上存在巨大的潛力。如今自然語言處理方法正得到迅速的發展，更多優于Word2Vec 的語言模型蘊含極大的能力，將其運用到從頭藥物設計的道路中，或許能使探索自動化藥物發現的前景更加可觀。