編輯丨ScienceAI
在藥物研發和材料科學領域,設計具有特定性質的分子是一項關鍵但耗時的任務。傳統方法依賴大量實驗試錯,而生成式 AI 模型的出現為這一領域帶來曙光。
然而,現有方法面臨兩大瓶頸:一是多數生成模型僅能處理連續型資料(如影像畫素),難以應對分子設計中離散的原子型別與連續座標共存的多模態資料;二是針對特定性質的定向生成通常需要額外訓練模型,成本高昂。
近日,來自北京大學、卡耐基梅隆大學、史丹佛大學等機構的研究團隊提出 TFG-Flow 框架,在無需額外訓練的條件下,成功實現了多模態資料的定向分子生成。該方法在量子屬性匹配、結構相似性最佳化、藥物分子設計等任務中展現出顯著優勢,相關論文已被 ICLR 2025 接收。
突破多模態引導的核心挑戰
分子可被視為由離散的原子型別(如碳、氧、氮)和連續的3D座標構成的多模態資料。現有生成模型在引導這類資料時面臨兩大難題:
1.維度災難:離散變數的組合空間隨原子數量呈指數增長,傳統方法計算代價高昂
2.幾何不變性:由連續變數刻畫的分子的 3D 結構需滿足旋轉、平移不變性,普通梯度引導會破壞這一特性
TFG-Flow 透過創新性的雙路徑設計攻克了這些難題(圖 1):
- 離散變數:採用蒙特卡洛重要性取樣,將計算複雜度從指數級降至對數級
- 連續變數:引入旋轉等變圖神經網路(EGNN),確保座標變換不影響分子性質
圖 1:TFG-Flow 透過離散取樣與連續梯度引導的雙路徑實現多模態生成
實驗表現:全方位超越基線模型
研究團隊在 QM9、GEOM-Drug 等分子資料集上進行了系統驗證:
量子屬性精準匹配
以極化率(α)、偶極矩(μ)等 6 種量子屬性為目標,TFG-Flow 的平均絕對誤差(MAE)相比當前最優的無訓練(training-free)引導方法降低 20.3%。如表 1 所示,其效能甚至逼近需要專門訓練的條件生成模型。
表 1:TFG-Flow 在量子屬性任務中顯著優於傳統方法
結構相似性最佳化
當以特定分子子結構為引導目標時,TFG-Flow 將生成分子與目標結構的Tanimoto 相似度提升了 76.8%(QM9 資料集)和 22.4%(GEOM-Drug 資料集),展現出精準的結構控制能力。
藥物分子設計實戰
在 CrossDocked2020 資料集上,TFG-Flow 生成的分子與蛋白質靶點的結合能(Vina Score)達到-7.65,優於主流方法 Pocket2Mol(-7.23)和 TargetDiff(-7.32),且分子合成可行性(SA Score)保持合理水平。
技術亮點:高效且通用
TFG-Flow 的核心優勢體現在兩大技術創新:
理論保障下的高效取樣
針對離散變數,研究團隊提出基於重要性取樣的近似方法(圖 2)。理論分析表明,提出的估計方法具有相合性等優良性質,且相比傳統方法能夠指數級降低計算量。實驗觀察發現這一方法僅需 16 次取樣即可高精度估計轉移機率。
圖 2:取樣次數 K=16 時已能達到穩定效能
即插即用的預訓練模型相容
不同於需要從頭訓練的傳統方法,TFG-Flow 可直接呼叫預訓練分子模型(如UniMol)作為目標預測器。實驗表明,使用 UniMol 時 TFG-Flow 的 MAE 進一步降低,與專業訓練模型 EEGSDE 的差距顯著縮小。得益於當前各類基座模型(foundation model)效能的增強,研究團隊認為 TFG-Flow 在未來將有更強的潛力。
開啟分子設計新正規化
這項工作的價值不僅在於技術突破,更在於其方法論創新:
- 降低成本:無需針對每個新目標重新訓練模型,計算資源消耗大大降低
- 擴充應用:框架可遷移至蛋白質設計、材料發現等需處理多模態資料的科學領域
- 促進協作:開原始碼允許化學家直接引入領域知識定義目標函式
研究團隊表示,未來將探索更大規模的基礎模型與 TFG-Flow 的結合,進一步釋放 AI 在科學發現中的潛力。正如審稿人評價:「這項工作為生成模型在科學計算中的應用樹立了新標杆。」
附:本文作者
林昊葦是北京大學人工智慧研究院二年級博士生,導師為梁一韜教授和馬劍竹教授,研究方向為生成模型及其在 AI4Science科 學發現中的應用。擔任 ICML、NeurIPS、ICLR、CVPR、ACL(ARR)等頂會審稿人,在 ICML、NeurIPS、ICML、Nature Communications、TPAMI 等頂會頂刊發表 14 篇論文。
黎善達是卡耐基梅隆大學機器學習系三年級的博士生,導師為 Yiming Yang 教授和 Ameet Talwalkar 教授。他透過理論與實證研究理解機器學習模型的規律與侷限,為更有效地運用機器學習方法提供洞察。同時,他研究深度學習方法在數學推理、程式碼生成、微分方程求解、自然科學等領域的應用。
論文地址:https://arxiv.org/abs/2501.14216
程式碼開源:https://github.com/linhaowei1/TFG-Flow