基於雙任務的端對端無模板反應預測模型

重新排版| X

本文旨在介紹西南交通大學楊燕/江永全團隊在《應用智慧》雜誌上發表的研究成果，其中的第一作者是胡昊哲，碩士生

論文連結：https://doi.org/10.1007/s10489-023-05048-8

#相關程式碼：https://github.com/AILBC/BiG2S

作者以目前無模板逆合成領域興起的圖至序列模型框架為基礎，進一步在同參數量規模下嘗試建構一類在單一模型中同時解決逆合成預測與正向反應預測任務的模型BiG2S（雙向圖至序列）

#同時，作者對主流逆合成資料集USPTO-50k 進行了初步分析，探討了模型在訓練過程中對不同SMILES 片段的預測難度差異以及模型在驗證集上Top-k 匹配率的波動情況，並針對這些問題引入了不平衡損失函數以及改進了模型集成和束搜索策略

在對三個主要的反應預測資料集進行測試時，透過對逆合成和正向反應預測任務進行測試，以及對上述模組進行全面的消融實驗證明，BiG2S能夠在適當的參數規模下以單一模型處理逆合成和正向反應預測任務。與已有的基於預訓練和資料增強的無模板方法相比，BiG2S的整體預測能力同樣出色

#研究的背景

##逆合成與正向合成是有機化學、電腦輔助合成規劃（CASP）以及電腦輔助藥物設計（CADD）領域的基礎性挑戰

進行內容改寫時，需要將原文改寫為中文，同時保持原始意思的不變

早期的逆合成規劃系統直接依賴領域專家預先編碼的反應規則，或是基於物理化學的計算，而隨著深度學習的快速發展。目前領域內的主流方法則是建構一個任務特異性的神經網路框架以從數據驅動的角度完成反應預測任務。其中，不依賴特定先驗化學知識的無模板法通過其類似於端到端機器翻譯的簡潔思路以及靈活性逐漸成為了領域內的主流發展方向之一。

目前，大多數無模板逆合成模型的輸入和輸出都是分子的 SMILES 字串，即採用了序列至序列（Seq2Seq）的流程。這種方法能夠很好地利用在自然語言處理領域內已有的模型框架，以及針對於SMILES 表示方法的成熟的資料處理流程

然而，由於SMILES作為一維字串序列無法很好地表徵和利用分子圖所包含的二維/三維結構訊息，因此在這個領域中逐漸出現了使用分子圖代替SMILES作為模型輸入的圖至序列（Graph2Seq）方法，或者將分子圖的附加結構信息嵌入到SMILES序列中的序列至序列方法。這兩種方法都能很好地利用分子圖的豐富結構特徵

基於此，本文以新興的圖至序列方法為基礎，在原先基於SMILES的模型對逆合成與正向反應在預測任務同時訓練的相關探索的基準上，進一步全面地探究對這類雙任務模型的構建與實驗，同時也初步地探索與分析了模型在訓練過程中所展現的難度不平衡以及Top-k匹配率波動的問題；在此基礎上建構的BiG2S模型能夠較好地處理主流資料集中的逆合成與正向反應預測任務，並在不使用資料增強的情況下取得與其他無模板逆合成模型一致的反應預測能力

整體框架需要進行重寫

BiG2S整體結構是端對端的編碼器-解碼器，如圖1所示。編碼器端採用局部定向訊息傳遞圖網路和融入圖結構偏移資訊的全域圖Transformer來產生最終的分子圖節點表徵。解碼器則使用標準的Transformer解碼器以自回歸的方式產生目標分子的SMILES序列

需要注意的是，為了同時學習逆合成和正向反應預測，解碼器端的輸入額外包含了不添加位置資訊的雙任務標籤。同時，解碼器端的歸一化層和最終的線性層都有兩套參數，分別用於學習逆合成任務和正向反應預測任務

圖1： BiG2S 整體框架圖

需要雙任務訓練框架

逆合成與正向反應預測是兩個相關的任務，其中逆合成任務以產物作為輸入和反應物作為目標輸出，而正向反應預測任務則相反。這兩個任務之間存在著緊密的聯繫，因為它們可以透過互換逆合成任務的輸入和目標輸出來轉換為正向反應預測任務

因此，一些基於SMILES的無模板模型已經嘗試透過將反向合成與正向反應預測作為訓練目標，來提高對化學反應的理解，並且取得了一定的效果。基於這個思路，作者進一步嘗試將雙任務訓練引入圖到序列模型中

具體而言，作者基於先前在其他方法上使用的參數共享策略，在解碼器的歸一化層和最終的線性層內建構了兩套任務特定的參數。而在其他模組中，兩類任務共享一套參數。同時，在輸入的分子圖節點和解碼器的初始輸入序列中額外新增了雙任務標籤。這樣，即使在控制整體模型規模的情況下，模型也能夠區分兩類任務並學習它們不同的資料分佈

需要訓練和推理最佳化

在在訓練過程中，作者進一步記錄並分析了模型在訓練過程中所反映出的兩類問題

首先，作者記錄了在USPTO-50k中不同SMILES字元的出現頻次以及其在訓練時對應的預測準確率，如圖2所示。在訓練過程中，對於在訓練集中佔比分別為0.4%和0.3%的S和Br，它們之間整體預測準確率的絕對差異達到了8%。這初步顯示了不同的分子結構/片段間預測的難度存在明顯的差異，由此，作者透過引入不平衡損失函數(如Focal Loss)來緩解此類問題，從而使模型能夠更加關注訓練時準確率較低的分子片段

圖2：在USPTO-50k訓練集中，不同SMILES字元的出現頻次以及其在訓練時的整體預測準確率

此外，作者也記錄了模型在訓練期間對驗證集的預測結果品質變化，如圖3所示。作者發現，在USPTO-50k資料集的中後期訓練階段，模型在驗證集上的Top-1準確率仍然在不斷提高，但在Top-3、Top-5和Top-10的預測品質方面出現了明顯下降

為了在提升模型Top-1 預測質量的同時保持模型前十位反應物生成結果的整體質量，我們額外構建了一類基於自定義評價指標的模型集成策略。具體來說，我們建立了一類儲存模型的隊列，同時根據預先定義的評估指標（如 Top-1 準確率，加權的 Top-k 準確率等）對存入的模型進行排序。在整個訓練過程中，我們動態地存入待選模型並自動產生基於隊列中前 3-5 位的整合模型，從而保留 Top-k 預測品質最高的模型。在推理階段，我們也基於新的框架重新建構了更重視搜尋廣度的束搜尋策略，以提升模型Top-k 產生結果的整體品質

##圖3：在USPTO-50k 資料集中訓練期間模型對驗證集的Top-k 匹配率以及Top-k 無效生成分子比例的變化曲線

需要進行雙任務實驗中的基準資料集

作者在逆合成任務與正向反應預測任務中進行了實驗，使用了包含5 萬、50 萬以及100 萬條化學反應數據的數據集USPTO-50k、USPTO-MIT 、USPTO-full。實驗中比較了雙任務模型和單任務模型的表現。根據圖4的測試結果顯示

在小規模資料集中，BiG2S基於雙任務訓練在逆合成任務中取得了領先的預測精度，同時也保持了較高的正向反應預測精度；然而在偏向正向反應預測的USPTO-MIT資料集以及大規模資料集USPTO-full中，由於模型整體參數量的限制，雙任務訓練後的模型表現出現了降低。儘管如此，從雙任務模型以幾乎一致的參數量與小幅度的反應預測能力降低( Top-k 準確率的絕對差值位於0.5% 左右)獲得了同時處理逆合成任務與正向反應預測任務的能力這個角度來看，BiG2S 模型已經達到了預期目標

#圖4：BiG2S 的雙任務模型與單任務模型在三個基準資料集上的實驗結果，其中上標b 表示採用單一任務模型分別完成兩類任務

重新分析消融實驗

作者通过消融实验进一步验证了新的束搜索算法以及采用不平衡损失后 BiG2S 在不同数据集中进行预测时的最佳温度超参数。这里的温度超参数是指 Softmax 中用于控制输出概率分布的温度参数 T。实验结果如图 5 和图 6 所示

在针对束搜索算法的实验中，可以观察到OpenNMT在搜索宽度扩大至3倍的同时搜索耗时仅扩大至1.74倍，而新的束搜索算法在Top-1精度与OpenNMT一致的情况下整体的搜索耗时扩大了1-2倍；但在Top-10预测结果的质量上，新的束搜索算法与OpenNMT相比具有至少3%的绝对精度优势以及2%的有效分子比例优势，可以说新的束搜索算法以搜索耗时为代价带来了明显提升了模型整体Top-k搜索结果的质量

在对温度超参数进行实验时，研究人员发现，在小规模数据集上使用较大的温度参数可以显著提高整体的 Top-k 预测精度。而在更大规模的数据集中，由于 BiG2S 模型规模不能完全适应所有反应数据，此时选择较小的温度参数往往有助于模型搜索

图 5：BiG2S 采用的束搜索算法与之前无模板逆合成模型常用的基于 OpenNMT 的束搜索算法在搜索结果质量与搜索耗时上的比较

图 6：使用不平衡损失后的 BiG2S 在不同数据集上采用不同温度超参数(T)进行搜索的 Top-k 预测准确率

研究的结论显示...

在本文中，作者提出了一种名为BiG2S的无模板反应预测模型，该模型可以同时处理逆合成任务和正向反应预测任务。通过采用适当的参数共享策略和额外的双任务标签，BiG2S能够以较小的参数量在不同规模的数据集上完成逆合成任务和反应预测任务，且其整体预测能力与主流模型相当

为了解决模型训练中不同 SMILES 字符预测难度不均衡和 Top-k 预测精度波动的问题，作者引入了不平衡损失、基于自定义评价指标的模型自动集成策略和基于新框架的束搜索算法来缓解这些问题

BiG2S 在三个不同规模的主流数据集上都表现出了很好的双任务预测能力，而进一步的消融实验也证明了额外引入的训练与推理策略的有效性

以上是基於雙任務的端對端無模板反應預測模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！