編輯丨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