今天給大家介紹一篇關于分子表征學習的文章。分子表征學習（MRL）旨在将分子嵌入到實向量空間中。然而，現有的基于SMILES（簡化分子線性輸入系統）或GNN（圖神經網絡）的MRL方法要麼以SMILES字元串作為輸入，難以編碼分子的結構資訊，要麼過度強調GNN結構的重要性，而忽視了其泛化能力。是以，作者提出使用化學反應來協助學習分子表征，其核心思想在于保持分子在嵌入空間中的化學反應的等價性，即強制讓每個化學方程式的反應物嵌入和生成物嵌入的總和相等，該限制在保持嵌入空間的有序性和提高分子嵌入的泛化能力中被證明是有效的。此外，該模型可以使用任何GNN作為分子編碼器，與GNN結構無關。實驗結果表明，這種方法在各種下遊任務中都達到了最佳性能，超過了最佳基線方法。
1.引言
在化學中，一般用IUPAC命名法、分子式、結構式、骨架式等形式來表示一個分子，然而這些形式最初是為人類而不是計算機設計的。為了便于計算機了解和利用分子，MRL被提出，MRL将分子映射到低維空間，表示為稠密向量。分子的學習向量（又稱嵌入）可用于廣泛的下遊任務，如化學反應預測、分子性質預測、分子結構預測等。
研究人員已經提出了許多MRL方法，其中大部分都将SMILES字元串作為輸入，盡管這些語言模型功能強大，但它們卻很難從SMILES中學習到分子的原始結構資訊，因為SMILES是分子結構的一維線性化。而GNN在處理分子表示中局限于設計新穎精緻的GNN結構，忽略了MRL的本質——泛化能力。這激勵了作者探索GNN架構之外的方法。
在本文中，作者使用化學反應來協助學習分子表征，提高其泛化能力。化學反應通常由化學反應方程式表示，反應物在左側，生成物在右側，例如乙酸和乙醇的費希爾酯化反應。作者的想法是保持分子在嵌入空間的等價性。更重要的是，當分子編碼器是以總和作為readout函數的GNN時，該模型能夠自動隐式地學習同一類别内一組化學反應的反應模闆，學習反應模闆是提高分子表征泛化能力的關鍵。
作者将該模型稱為MoLR（chemical-reaction-aware molecule embeddings，化學反應感覺的分子嵌入）。并且将可視化分子嵌入，表明它們能夠編碼反應模闆和幾個關鍵的分子屬性，如分子大小和最小環的數量。
2.方法
結構分子編碼器
分子圖用G = (V, E)表示，其中V = {a1, …}是非氫原子的集合，E = {b1, …}是鍵的集合。在這項工作中，每個原子ai都有初始特征向量，該向量包括四個one-hot向量，每個one-hot向量表示一種原子屬性，四個one-hot向量分别表示四種原子屬性，分别為元素類型、電荷、原子是否是芳香環以及附着的氫原子數，然後每個one-hot向量有一個額外的未知條目，以在推理過程中處理未知值。此外，每個鍵bi都有一種鍵類型，例如單鍵、雙鍵，根據作者的實驗，鍵的類型不會持續改善模型性能，是以沒有将鍵的類型作為輸入。
作者使用GNN作為基礎模型，典型的GNN遵循領域聚合政策，通過聚合原子的鄰居和自身的表示來疊代更新原子的表示。GNN的第ｋ層是：
這裡，是原子ai在第k層的表征向量，N(i)是與ai相連的原子集合，K是GNN的層數，這個聚合函數的選擇是設計GNN的關鍵。
最後，readout函數用于聚合最後一個GNN層輸出的所有結點表示，以獲得整個分子的表示hG：
readout函數可以是簡單的置換不變函數，例如求和和取平均，也可以是更複雜的graph-level pooling算法。
圖1：(a) GNN編碼器處理脯氨酸分子，省略了氫原子；(b) 丙酸和丙醇的費希爾酯化反應，以及模型學習到的相應反應模闆。反應中心用橙色表示，距離反應中心1或2的原子用淺橙色表示；(c)小批量化學反應對比損失的圖示，dij是嵌入和之間的歐氏距離。
保持化學反應等價性
如前所述，作者希望在分子嵌入空間中保持反應物和生成物的等價性，即：
其中，R是反應物集合，P是生成物集合。這個簡單的限制能有效提高分子嵌入的品質，下面的命題将會證明化學反應中的等價關系是等式(3)限制下的等價關系：
對于一個分子的集合M，M的所有子集2的M次方可以根據等價關系劃分為等價類，一個等價類中所包含的所有分子的嵌入總和應該相等。是以，等式3形成了一個線性方程組，化學反應等價性對基礎分子的嵌入施加了更加強大的限制，分子嵌入的可行解将更加穩健，整個嵌入空間将更加有序。
然後作者提出反應中心的概念，反應中心被定義為将反應物轉換為生成物所需的一組最小圖形編輯，鑒于反應中心的概念，有以下命題：
命題2表明，反應物嵌入和生成物嵌入之間的剩餘僅取決于距離反應中心小于K跳的原子。如圖1 (b)所示，剩餘物完全取決于反應中心（橙色）以及距離反應中心1或2的原子（淺橙色）。這意味着，方程
适用于任何官能團R1和R2，這稱為反應模闆，它概括了同一類别内的一組化學反應，反應模闆可以輕松地适應訓練集中沒見過但符合已知反應類型的反應。
訓練模型
在這裡作者使用了一個類似于基于小批量的對比學習架構。對于小批量的資料B = {R1→P1,R2→P2, …}，首先使用GNN編碼器處理B中所有的反應物和生成物，并得到它們的嵌入物，比對的反應物-生成物對(Ri, Pi)被稱為正對，其嵌入差異将最小化，而不比對的反應物-生成物對(Ri, Pj)(i≠j)，其嵌入差異将最大化。為了避免總損失被負對主導，作者使用基于邊距的損失：
其中γ是一個邊距超參數。是以，可以使用基于梯度的優化方法（如随機梯度下降SGD）最小化上述損失來訓練模型。
3.實驗
化學反應預測
作者使用USPTO-479k資料集，每個反應執行個體包含5種反應物和1種生成物。化學反應預測的任務可表述為一個排序問題。在推斷階段，給定化學反應的反應物集R，将測試集中的所有生成物視為候選者C，根據反應物嵌入hR和候選生成物嵌入hC之間的L2距離對所有生成物進行排序，和真實值比較以計算名額。結果如表1所示，很明顯，MoLR的方法都顯著優于基線方法。
表1：USPTO-479k資料集上的化學反應預測結果。最佳結果以粗體突出顯示，基線的最佳結果以下劃線突出顯示。
案例研究 作者選取了前20個反應作為案例研究，表2展示了兩個反應的結果，結果表明MoLR模型在預測生成物時非常精确。
表2：在USPTO-479k資料集上的案例研究
現實場景生成物預測中的多選擇問題 每個問題給出反應的反應物，從4或5個選項中選出正确的生成物，結果如圖2所示，MoLR優于基線方法。
圖2：現實場景生成物預測中的多選擇問題的結果
分子性質預測
作者在5種資料集中測試MoLR，每個資料集包含數千個SMILES分子。預測的AUC結果如表3所示，MoLR在4個資料集中表現最好，作者将MoLR在分子性質預測方面的優異性能歸因于，MoLR在USPTO-479k上預訓練，是以根據命題2對反應中心敏感。注意，反應中心通常由化學活性官能團組成，這些官能團對确定分子性質至關重要。
圖編輯距離預測
表4：分子性質預測的AUC結果
圖編輯距離（GED）是兩個圖之間相似性的度量，定義為将一個圖轉換為另一個圖的最小編輯操作數。這項任務根據兩個分子圖的嵌入來預測它們之間的GED，目的是顯示學習到的分子嵌入是否能夠保持分子之間的結構相似性。此外，計算GED是NP-hard問題，這項任務也能近似作為計算GED的算法，結果如表4所示。
表5：在QM9資料集上GED預測的RMSE結果，最佳結果加粗顯示
嵌入可視化
圖4：在BBBP資料集上可視化分子嵌入空間
在圖4a中，分子根據滲透性的性質着色，我們發現了兩個非滲透性分子群落，這表明molr可以捕獲感興趣的分子性質。在圖4b中，根據BBBP資料集中随機選擇的分子（編号1196）的GED對分子進行着色，顯然，與1196号分子結構相似的分子（橙色）在嵌入空間中很接近，而與1196号分子結構不同的分子（紅色）在嵌入空間中相距很遠，這表明molr能捕捉分子的結構相似性。在圖4c中，分子根據其大小着色，嵌入空間被完美地分割為小分子區域（上部）和大分子區域（下部）。
換句話說，二維嵌入空間的縱軸表征分子大小。驚訝的是，我們發現橫軸與分子中最小環的數量有關：如圖4d所示，沒有環的分子（藍色）僅存在于左簇中，含有一個環的分子（黃色）僅存在于左簇和中間簇中，含有兩個環的分子（橙色）基本上位于中間簇，而右邊簇主要由含有3個及以上環的分子（紅色）組成。
圖3：醇氧化和醛氧化反應的可視化
作者還以醇氧化和醛氧化為例說明MoLR編碼化學反應，結果如圖3所示，藍色箭頭的長度大約是相應的紅色或橙色箭頭的兩倍，這正是因為(CH2OH)2/(CH2CHO)2具有兩個待氧化的羟基/醛基。
4.結論和未來工作
在這項工作中，作者使用GNNs作為分子編碼器，并使用化學反應來協助學習分子表征，強制讓反應物的嵌入總和等于生成物的嵌入總和。該模型能夠學習反應模闆，這是提高模型泛化能力的關鍵。該模型能夠勝任廣泛的下遊任務，可視化的結果表明，學習到的嵌入是有組織的和反應感覺的。
作者指出了今後工作的四個方向。首先，環境條件也是化學反應的一部分，需要考慮進去；第二，值得研究如何顯式輸出學習的反應模闆；第三，值得研究如何在嵌入空間中區分立體異構體，因為現有模型不能處理立體異構體；最後，加入額外的資訊（如分子的文本描述）來協助學習分子表征也是一個很有希望的方向。