作者 | 香港理工大學李佳潼
編輯 | ScienceAI
去年，香港理工大學研究團隊開發了一個基於檢索的提示正規化MolReGPT，利用大型語言模型探索分子發現，在分子和分子文字描述之間進行翻譯。
近日，香港理工大學、上海交通大學和上海人工智慧實驗室聯合發表了題為《Large Language Models are In-Context Molecule Learners》的文章，也是MolReGPT[1]的續作。
論文連結：https://arxiv.org/abs/2403.04197
模型連結：https://huggingface.co/phenixace/
介紹
近來，大語言模型在生物化學領域展現出了優異的效能，尤其是分子-描述翻譯任務[1]，這減小了自然語言描述與分子空間的差異。
但是，之前的方法要麼是需要進行大量的領域預訓練，要麼是在分子與描述的對齊上存在過於粗糙的問題[2]，要麼是對大語言模型的能力和引數規模有著嚴苛的要求[1]。
為了解決這些問題，我們提出了上下文分子微調（ICMA），作為一個讓大語言模型學習分子-描述對齊的新方法。
具體來說，上下文分子微調由以下三個部分組成：混合模態上下文召回，召回後重排，和上下文分子微調。
- 首先，混合模態上下文召回沿用了BM25和分子圖召回分別對分子描述和分子進行召回，以增強模型的輸入。
- 此外，為了解決召回中存在的不準確、樣本間過度重複等問題，我們設計了兩個召回後處理方法：序列反轉和隨機遊走，以提升召回結果的質量。
- 最後，上下文分子微調藉助於大語言模型的上下文學習能力，生成最終的分子/分子描述。
我們在ChEBI-20[3]和PubChem324K[2]兩個分子-分子描述資料集上展開實驗，實驗結果表明，上下文分子微調可以使得Mistral-7B在分子-分子描述上取得SOTA或接近的結果，無需領域預訓練和複雜的模型結構。
我們的貢獻在於：上下文分子微調不需要額外的領域預訓練就可以發揮作用，為最新最先進的十億甚至百億引數大語言模型在化學任務上的運用提供了新思路。與此同時，上下文分子微調透過學習上下文例子中分子描述片段與分子SMILES結構的對應關係，精細化了分子描述和分子的對齊。此外，上下文分子微調不那麼依賴於模型的上下文學習和推理能力，即便是稍小的模型也能獲得良好的效能增益。
方法
分子的相似性原理指出，相似的分子一般會具有相似的性質，如圖1所示。藉助於分子SMILES和分子描述片段的對應，我們可以基於這些共現資訊大致推理出，如果有另一個相似的分子，它的SMILES片段會對應哪些相應分子描述片段，即具有什麼樣的結構和化學性質。
圖1: 三個相似分子以及他們的分子描述。分子可以被表示為SMILES表示式或者分子圖，而分子描述說明了分子的特徵。這裡三個分子在圖結構上是相似的，分子描述中重疊的地方被藍色和粉色標出。
因此，在這篇文章中，我們提出了上下文分子微調（ICMA），其流程如圖2所示，總共三個階段：混合模態上下文召回，召回後重排，和上下文分子微調。
圖2: 上下文分子微調（ICMA）框架。
混合模態上下文召回
召回的質量對上下文例子的資訊量至關重要，如果我們召回到的分子足夠相似，那麼他們在描述中的重疊就很可能會更多，這樣模型就可以更好地學習到分子SMILES片段與描述片段的對應。
本文提出了混合模態上下文召回，在分子描述召回過程中，沿用了BM25對文字模態進行召回。而在分子召回過程中，我們提出了分子圖召回，從圖模態來召回相似的分子。
在分子到描述任務中，ICMA首先採用分子圖召回（取代了MolReGPT[1]的Morgan分子指紋）去提升分子召回的質量。分子圖召回首先採用一個圖神經網路去對分子圖進行編碼，得到分子圖向量:
召回後重排
上下文分子微調
研究結果
本文的實驗在ChEBI-20和PubChem324K兩個資料集上展開。表2、3列出了ChEBI-20資料集上，上下文分子微調和直接微調、基線模型的效果比較。首先，在兩個子任務上，上下文分子微調相較於直接微調都可以帶來可觀的效能增益，而相較於先前的基線模型，上下文分子微調也可以取得SOTA或者相接近的效能。
表4、5列出了PubChem324K資料集上，上下文分子微調與直接微調的效果比較，其結果與ChEBI-20的結果相類似，上下文分子微調相較於直接微調都帶來了可觀的表現增益。
與此同時，本文也展開了大量的實驗去探究影響上下文分子微調效果的因素，集中探究了召回演算法對效果的影響、上下文設定的影響、不同基座模型的影響（模型型別和模型引數量）
表6、7展示了使用不同召回演算法進行上下文分子微調的效果區別，顯示出BM25和分子圖召回所帶來的效果增益。
圖3則說明了上下文設定（例子數量和截斷長度）對上下文分子微調的影響，表明截斷長度必須適應上下文長度，以及例子增多所可能帶來的負面影響（如過多的噪聲）。
圖3: 模型效能隨著上下文設定（例子數量和截斷長度）變化的比較。
而圖4則展現了上下文分子微調所具備的scaling law，更大的模型因其具有更強的上下文學習和推理能力，因而會獲得更好的上下文分子微調效果。
圖4: 上下文分子微調（ICMA）的scaling law，這裡選取了Galactica-125M，Galactica-1.3B，和Mistral-7B進行比較。
表11、12進一步說明了上下文分子微調的通用性，可以使得不同的大模型都能夠在分子-描述翻譯任務上取得效能增益，而不僅僅侷限於某一個或者某一類模型。
最後，如表8、9所示，消融實驗也論證了召回後重排策略所帶來的效能增益。
結論
本文提出了上下文分子微調，作為一個新的正規化，來將大語言模型運用到分子-描述翻譯任務上，不需要額外的領域預訓練和複雜的模型結構，上下文分子微調使得大語言模型能夠運用他們的上下文學習和推理能力學習到分子和描述之間的對齊，提升了大語言模型在分子-描述翻譯任務上的能力，也表明了大語言模型是天然的上下文分子學習者。我們的方法為將十億以上的大語言模型運用到科學領域提供了新方法和新思路。
參考文獻：
[1] Li, J., Liu, Y., Fan, W., Wei, X. Y., Liu, H., Tang, J., & Li, Q. (2024). Empowering molecule discovery for molecule-caption translation with large language models: A chatgpt perspective. IEEE Transactions on Knowledge and Data Engineering.
[2] Liu, Z., Li, S., Luo, Y., Fei, H., Cao, Y., Kawaguchi, K., ... & Chua, T. S. (2023). Molca: Molecular graph-language modeling with cross-modal projector and uni-modal adapter. arXiv preprint arXiv:2310.12798.
[3] Edwards, C., Lai, T., Ros, K., Honke, G., Cho, K., & Ji, H. (2022). Translation between molecules and natural language. arXiv preprint arXiv:2204.11817.
[4] Xia, J., Zhao, C., Hu, B., Gao, Z., Tan, C., Liu, Y., ... & Li, S. Z. (2022, September). Mole-bert: Rethinking pre-training graph neural networks for molecules. In The Eleventh International Conference on Learning Representations.