編輯 | 蘿卜皮

眾所周知，藥物研發需要經歷一段極其漫長且復雜的過程。好在大型語言模型 (LLM) 在藥物設計中的應用，正逐漸給這一挑戰帶來曙光。不過，現有基于 LLM 的方法往往難以有效地整合三維分子結構。

在最新的研究中，浙江大學侯廷軍團隊提出了 Token-Mol，這是一個僅包含 token 的 3D 藥物設計模型，它將二維和三維結構信息以及分子屬性編碼成離散的 token。

Token-Mol 建立在 Transformer 解碼器之上，并經過因果掩蔽訓練，引入了專為回歸任務定制的高斯交叉熵損失函數，從而在多個下游應用中實現了卓越的性能。

能力方面，該模型超越了現有方法，在兩個數據集上將分子構象生成性能分別提高了 10% 和 20% 以上，同時在性質預測方面比僅使用 token 的模型高出 30%。

在基于口袋的分子生成中，它分別將藥物相似性和合成可及性提高了約 11% 和 14%。并且，Token-Mol 的運行速度比「專家」擴散模型快 35 倍。

該研究以「Token-Mol 1.0: tokenized drug design with large language models」為題，于 2025 年 5 月 13 日發布在《Nature Communications》。

背景與挑戰

在藥物研發中，獲取帶注釋數據集的高成本仍然是該領域發展的重大障礙。近年來，以 BERT 和 GPT 為代表的無監督學習框架的快速發展，為化學和生物等學科引入了無監督化學和生物預訓練模型。

這些模型經過大規模無監督訓練，學習小分子或蛋白質的表征，并隨后針對特定應用進行微調。通過在大規模數據集上進行無監督學習，這些預訓練模型有效地解決了與稀疏標記和分布外泛化能力欠佳相關的挑戰。

大規模分子預訓練模型大致可分為兩大類：基于化學語言的模型和利用分子圖的模型。

化學語言模型使用諸如簡化分子輸入系統 (SMILES) 或自引用嵌入字符串 (SELFIES) 等表示形式對分子結構進行編碼。它們往往采用類似于 BERT 或 GPT 的訓練方法。

基于圖的分子預訓練模型表現出更高的通用性。它們以圖形形式表示分子，節點表示原子，邊表示化學鍵。預訓練方法包括各種技術。與基于語言的模型不同，基于圖的分子預訓練模型本質上包含幾何信息。

但是這里有一個很尷尬的問題，基于化學語言的大規模分子預訓練模型無法處理對分子特性至關重要的3D結構信息，所以它難以勝任相關下游任務；基于圖結構的預訓練模型雖能整合 3D 信息卻僅聚焦性質預測，同時難以與通用 NLP 模型融合。

因此，當前亟需開發能克服這兩類模型局限、適用于所有藥物設計場景并可無縫對接通用大語言模型的新型預訓練模型。

Token-Mol

在本研究中，為了克服現有模型的局限性，浙江大學侯廷軍、康玉、Chang-Yu Hsieh 團隊提出了 Token-Mol，一個用于分子預訓練的大規模語言模型。

為了增強與現有通用模型的兼容性，研究人員采用了僅使用 token 的訓練范式，將所有回歸任務重新定義為概率預測任務。Token-Mol 采用 Transformer 解碼器架構構建，通過 SMILES 和扭轉角 token 整合必要的二維和三維結構信息。

具體來講，Token-Mol 的完整工作流程如下圖所示。初始階段通過隨機因果掩蔽對數據集進行預訓練。隨后，模型在定制數據集上進行微調，以完成特定的下游任務，包括構象生成、基于口袋的分子生成以及多種屬性的預測。對于回歸任務，在微調過程中使用 GCE 損失函數。此外，還可以使用強化學習進一步優化特定下游任務的性能。

圖示：Token-Mol 概述。（來源：論文）

在此背景下，預訓練數據集的預處理至關重要。對標準 SMILES 格式的整個分子進行深度優先搜索（DFS）遍歷，以提取分子結構中嵌入的扭轉角。之后，每個提取出的扭轉角都會被同化為一個 token，并附加到 SMILES 字符串中。

在整個預訓練階段，研究人員會實施基于因果回歸的隨機因果掩蔽。預訓練完成后，研究人員會對下游任務進行微調。值得注意的是，任務提示是專門為構建對話系統而設計的。

這一特性顯示了僅使用 token 的模型相較于其他大規模模型的一個關鍵優勢：它們能夠促進實時交互。

對于基于口袋的分子生成任務，研究人員引入了口袋編碼器和融合模塊，從而更好地將蛋白質口袋信息整合到模型中。同時，他們利用多頭條件注意力機制，將每個自回歸步驟中生成的信息完整地整合到后續迭代中。該機制將自回歸過程中生成的每個 token 視為后續生成的先決條件，從而確保整個查詢、鍵和值矩陣均源自原始序列。

需要注意的是，在實際應用中，先導化合物不僅要對靶標表現出高親和力，還要滿足一系列標準，包括高生物活性和多種良好的藥理特性。這對基于口袋的分子設計任務提出了更高的要求，而訓練數據集中受體-配體分子對的整合存在固有的局限性。

該模型主要利用來自蛋白質口袋的信息來生成配體分子。因此，這些生成分子的性質受到訓練數據的嚴重影響，限制了對其生物物理和化學性質的明確控制。當需要精確調節分子特性時，這些限制尤為明顯。

Token-Mol 建立在自回歸語言模型架構之上，其中 token 的生成與強化學習框架中的動作相一致，從而促進了強化學習的無縫優化，從而保證獲得量身定制的結果。

性能驗證

為了驗證 Token-Mol 的性能，研究人員將它應用于分子構象生成、性質預測和基于口袋的分子生成等任務。

圖示：Token-Mol 和基線模型之間的分子特性分布。（來源：論文）

在分子構象生成任務中，Token-Mol 的表現優于現有的最佳方法，在兩個數據集上的各項指標上分別實現了超過 10% 和 20% 的準確率提升。

在分子屬性預測任務中，Token-Mol 與僅使用 token 的模型相比，在回歸任務中平均提升了 30%。

在基于 Pocket 的分子生成任務中，Token-Mol 不僅生成了與當前最佳模型生成的分子 Vina 得分相當的分子，而且藥物相似性 (QED) 和合成可及性 (SA) 分別提升了約 11% 和 14%。

圖示：真實藥物設計場景評估。（來源：論文）

為了進一步驗證 Token-Mol 的泛化能力，團隊還在真實的藥物設計場景中進行了測試。值得一提的是，Token-Mol 生成的類藥分子在 8 個不同靶點的評估中，平均成功率提高了 1 倍。

此外，將強化學習集成到 Token-Mol 中，能夠有效提升特定下游任務在更現實場景下的性能，使其在 3D 任務上優于基于幾何圖神經網絡的大型模型。

與 Token-Mol 聊天

Token-Mol 的 token-only 框架相較于傳統回歸模型具有顯著優勢，因為它能夠無縫集成前沿的大模型技術，包括快速學習、混合專家（MoE）和檢索增強生成（RAG）。

為了說明此功能，該團隊在論文里展示了幾個簡單的對話用例。通過運用提示學習，研究人員可以控制諸如本研究中提到的屬性預測等任務的執行。

首先，他們插入特定的提示，例如「預測 ESOL」從而微調模型。這使得用戶能夠在提示后與模型直接交互，從而請求預測不同的分子屬性。

圖示：與 Token-Mol 聊天。（來源：論文）

使用時，研究人員查詢了不同分子的各種屬性，Token-Mol 成功地給出了相應的預測。這展現了 Token-Mol 擁有與化學家進行有意義對話的潛力。

用戶可以提供分子構象，但由于 Token-Mol 可以生成相應的構象，因此最終輸出將僅包含預測的目標屬性。

研究人員表示，未來該模型可以整合 RAG。

當向 Token-Mol 查詢分子的特定屬性時，系統會使用基于嵌入的向量搜索，將查詢轉換為向量。然后，該向量會與數據庫中高度相關的向量描述進行匹配，以提供上下文信息。查詢以及檢索到的上下文信息（例如空間結構信息和其他相關屬性）隨后被輸入到 Token-Mol，最終生成答案。

總而言之，Token-Mol 的內在潛力巨大。它的新穎架構也為藥物設計 AI 模型的標準化提供了參考。

論文鏈接：https://www.nature.com/articles/s41467-025-59628-y