在當今科學研究與工業應用的前沿領域,原子系統三維結構的生成建模正展現出顛覆性潛力,有望徹底重塑新型分子和材料的逆向設計版圖。從精準的結構預測到靈活的條件生成,當前最先進的擴散模型及流匹配模型已在生物分子解析、新材料研發及基於結構的藥物設計等關鍵任務中嶄露頭角,成為科研人員突破技術瓶頸的核心工具。
然而,在這一蓬勃發展的領域背後,一個關鍵難題始終制約着技術躍遷——現有模型缺乏跨系統的通用性。 儘管所有原子系統均遵循相同的物理原理決定其三維結構與相互作用,但小分子、生物分子、晶體及其複合體系的建模卻長期處於「分而治之」的狀態。絕大多數擴散模型高度依賴特定系統的固有特性,需在分類數據(如原子類型)與連續數據(如三維座標)交織的複雜乘積流形上進行多模態生成,這導致不同系統間的模型難以兼容。
以具體場景為例:小分子的從頭生成需拆分為原子類型(分類)與三維座標(連續)兩個獨立擴散過程,去噪模型雖需學習二者的協同演化規律,卻常因中間狀態失真降低採樣效率;生物分子建模需額外引入旋轉流形,將原子組視為剛體處理;而晶體與材料的擴散過程則必須兼容週期性特徵,在原子類型、分數座標、晶格參數等多維度參數構成的聯合流形上運行 —— 這些差異使得跨系統統一建模成為領域內長期未決的挑戰。
在此背景下,Meta 基礎人工智能研究(FAIR)、劍橋大學與麻省理工學院的聯合科研團隊提出了突破性解決方案——全原子擴散 Transformer(All-atom Diffusion Transformer,簡稱 ADiT)。
作為一種基於 Transformer 的統一潛在擴散框架,ADiT 的核心優勢在於打破了週期性與非週期性系統之間的建模壁壘,通過全原子統一潛在表示與 Transformer 潛在擴散兩大創新,實現了用單一模型生成分子與晶體。 其設計幾乎不引入歸納偏差,使得自編碼器與擴散模型在訓練和推理效率上遠超傳統等變擴散模型 —— 在相同硬件條件下,生成 10,000 個樣本的時間從 2.5 小時縮短至 20 分鐘以內。更值得關注的是,當模型參數擴展至 5 億規模時,其性能呈現可預測的線性提升,這一特性為構建通用型生成化學基礎模型奠定了關鍵基礎,標誌着原子系統建模向通用性與規模化應用邁出了里程碑式的一步。
相關研究成果以「All-atom Diffusion Transformers: Unified generative modelling of molecules and materials」為題,入選 ICML 2025。
研究亮點:
- ADiT 首次實現了適用於週期性材料和非週期性分子系統生成模型的統一
- ADiT 依託全原子統一潛在表示與使用 Transformer 進行潛在擴散,有效簡化了生成過程,且幾乎無歸納偏差
- ADiT 具有出色的可擴展性與效率,訓練和推理速度遠超等變擴散模型
論文地址:
https://go.hyper.ai/27d7U
更多 AI 前沿論文:
https://go.hyper.ai/owxf6
數據集:從週期到非週期,覆蓋多領域實驗數據
在該研究中,研究團隊首先選取了多類具有代表性的數據集展開實驗:
- MP20 數據集, 包含 45,231 個來自 Materials Project 的亞穩態晶體結構,單元格內最多有 20 個原子,涵蓋 89 種不同元素,能很好地代表週期性材料系統;
- QM9 數據集, 由 130,000 個穩定有機小分子構成,最多含 9 個重原子(C、N、O、F)及氫原子,是典型的非週期性分子系統代表;
- GEOM-DRUGS 數據集, 包含 430,000 個最多 180 個原子的大型有機分子;
- QMOF 數據集, 含有 14,000 個金屬有機框架結構。
其中,MP20 與 QM9 分別對應不同類型的原子系統, 為模型在週期性與非週期性系統上的聯合訓練提供了基礎,且研究團隊按先前研究的方式劃分數據,保證了與其他模型比較的公平性;GEOM-DRUGS 與 QMOF 進一步拓展了模型測試的範圍,能更全面地檢驗模型的泛化能力。
ADiT:以雙核心思想構建統一原子系統生成模型
ADiT 作為一種潛在擴散模型,其核心設計圍繞兩個關鍵思想展開,以此實現對週期性和非週期性原子系統的統一生成建模。
第一個關鍵思想是全原子統一潛在表示, 研究團隊將週期性與非週期性原子系統均視為三維空間中的原子集合,進而開發出一種包含每個原子分類屬性(如原子類型)和連續屬性(如三維座標)的統一表示方式。通過訓練變分自編碼器(VAE)進行全原子重建,該編碼器能夠將分子和晶體嵌入到一個共享的潛在空間中, 這就為不同類型原子系統的統一處理搭建了基礎框架。
第二個關鍵思想是使用 Transformer 進行潛在擴散, 在 VAE 編碼器構建的潛在空間中,研究團隊引入擴散 Transformer(DiT)開展生成建模工作。在推理過程中,藉助無分類器引導技術,能夠對新的潛在變量進行採樣, 而這些潛在變量又可以通過 VAE 解碼器重建為有效的分子或晶體,從而完成從潛在空間到實際原子系統的轉化。
基於這兩個核心思想,ADiT 的實驗方法分為兩個階段有序推進。
第一階段,研究人員構建用於重建的自動編碼器(Autoencoder for reconstruction), 通過 VAE 對分子和材料的全原子表示進行聯合重建,學習並構建共享的潛在空間——這是實現不同原子系統統一建模的前提,為後續生成過程奠定基礎。
第二階段,研究人員構建潛在擴散生成模型(Latent diffusion generative model), 利用 DiT 從潛在空間生成新樣本,這些樣本經無分類器引導解碼為有效的分子或晶體。這種潛在擴散設計的顯著優勢在於,將處理分類與連續屬性的複雜性轉移至自編碼器,使潛在空間中的生成過程更簡化且可擴展,有效提升了模型處理不同原子系統的效率與適應性。
ADiT 分兩個階段對化學系統進行生成建模
ADiT 在晶體與分子生成中性能領先
為充分凸顯 ADiT 的性能優勢,研究團隊選取了多類基線模型進行針對性對比。在晶體生成領域, 對比對象包括 CDVAE、DiffCSP、FlowMM 等基於多模態乘積流形的等變擴散和流匹配模型,以及非等變擴散模型 UniMat、兩階段框架 FlowLLM;在分子生成領域, 則與等變擴散模型、GeoLDM、Symphony 等模型展開比較。通過與這些領域內先進基線模型的系統對比,ADiT 的性能優勢得到了清晰展現。
從具體實驗結果來看,ADiT 在晶體和分子生成任務中均達到了 SOTA 水平。 在晶體生成方面,ADiT 生成的晶體在有效性、穩定性、獨特性和新穎性等關鍵指標上表現出色。在分子生成任務中,ADiT 在 10,000 個採樣分子的有效性和獨特性指標上位居前列。
ADiT 的聯合訓練機制也為其帶來了顯著的性能增益。實驗數據顯示,同時在 QM9 和 MP20數據集上訓練的 ADiT,在材料和分子生成任務中全面優於僅在單一數據集上訓練的版本。
模型規模的擴展對 ADiT 的性能提升具有可預測性。如下圖所示,隨着 DiT 去噪器參數量從 3,200 萬(ADiT-S,藍色)增至1.3 億(ADiT-B,橙色),再到 4.5 億(ADiT-L,綠色),即使在約 13 萬個樣本的中等規模數據集上,擴散訓練損失持續降低,有效性比率穩步提升,呈現出顯著的規模效應。這種模型規模與性能的強相關性表明,通過擴大模型參數和數據量,有望推動 ADiT 實現進一步突破。
ADiT 去噪參數數量增加對訓練損失和生成有效性的影響
在效率方面,ADiT 相比等變擴散模型展現出顯著的速度優勢。 如下圖所示,在英偉達 V100 GPU 上生成 10,000 個樣本時,基於標準 Transformer 的 ADiT 在積分步數上的擴展性遠超使用計算密集型等變網絡的 FlowMM 和 GeoLDM。即便 ADiT-B 的參數規模比等變基線大 100 倍,其推理速度仍更快,這凸顯了 Transformer 架構在擴展實用性上的優勢。
ADiTs 和等變擴散模型生成 10,000 個樣本的時間關係圖
此外,ADiT 在更大系統上的擴展性也得到了驗證。在包含 43 萬個最多 180 個原子的 GEOM-DRUGS 分子數據集上,ADiT 與最先進的等變擴散和流匹配模型相比,在有效性和 PoseBusters 指標上表現相當。 值得注意的是,ADiT 基於標準 Transformer 架構,幾乎不引入分子歸納偏差,且無需顯式預測原子鍵,卻能達到與等變模型相當的性能,進一步體現了其設計的通用性和廣泛適用性。
產研共同推進原子系統三維結構生成的突破創新
事實上,在原子系統三維結構生成建模這一前沿研究領域,學術界與企業界歷經不懈探索,已有諸多成果備受矚目。
在學術界,加州大學伯克利分校、微軟研究院攜手 Genentech 公司的研究團隊推出多模態蛋白質生成方法 PLAID, 該方法巧妙藉助預訓練權重中的結構信息,以 DiT 執行去噪任務,於不同蛋白質長度的結構質量與多樣性分析中,較其他基準方法彰顯出更為卓越的性能。
企業界在這一領域同樣踴躍投身探索,創新驅動發展。中國生成式 AI 蛋白質設計創新企業百奧幾何重磅發佈的全球首款全能蛋白質基礎模型 GeoFlow V2, 構建起統一的原子擴散模型架構,一舉攻克蛋白質結構預測與設計任務。在抗體及抗原 - 抗體複合物結構預測方面,GeoFlow V2 憑藉其超凡的精度與速度,全面領先同類產品。字節跳動推出的 Seedance 1.0 則另闢蹊徑,採用變分自編碼器與擴散 Transformer 相結合的技術方案,實現了快速高效的 AI 視頻生成,其速度優勢為實時創作和交互式應用打開了全新局面,預示着在商業應用領域擁有廣闊前景。
這些學術界的科研突破與企業界的創新實踐,共同推動着原子系統三維結構生成建模領域的發展。隨着技術的持續進步,這一領域必將在新型材料研發、藥物設計等諸多方面發揮更大作用,為解決全球性的科學難題和產業挑戰提供強有力的支撐。
參考文章:
1.https://mp.weixin.qq.com/s/oF3-y7z8u1XpEtjd4q1u4w
2.https://mp.weixin.qq.com/s/tK0-