Uni-3DAR用自回歸統(tǒng)一微觀與宏觀的3D世界,性能超擴(kuò)散模型256%,推理快21.8倍
從微觀世界的分子與材料結(jié)構(gòu)、到宏觀世界的幾何與空間智能,創(chuàng)建和理解 3D 結(jié)構(gòu)是推進(jìn)科學(xué)研究的重要基石。3D 結(jié)構(gòu)不僅承載著豐富的物理與化學(xué)信息,也可為科學(xué)家提供解構(gòu)復(fù)雜系統(tǒng)、進(jìn)行模擬預(yù)測和跨學(xué)科創(chuàng)新的重要工具。如何準(zhǔn)確且高效地構(gòu)建 3D 模型、理解和生成 3D 世界正在成為 AGI、AI for Science、具身智能三大 AI 熱門領(lǐng)域共同關(guān)注的焦點。而隨著 AI 技術(shù)的發(fā)展,大型語言模型(LLM)與大型多模態(tài)模型(LMM)那強(qiáng)大的自回歸下一 token 預(yù)測能力也已經(jīng)在開始被用于創(chuàng)建和理解 3D 結(jié)構(gòu)。基于此,我們看到了 AI for Science 的新可能。
近日,一個開創(chuàng)性的此類大模型誕生了!
它名為 Uni-3DAR,來自深勢科技、北京科學(xué)智能研究院及北京大學(xué),是一個通過自回歸下一 token 預(yù)測任務(wù)將 3D 結(jié)構(gòu)的生成與理解統(tǒng)一起來的框架。據(jù)了解,Uni-3DAR 是世界首個此類科學(xué)大模型。并且其作者陣容非常強(qiáng)大,包括了深勢科技 AI 算法負(fù)責(zé)人柯國霖、中國科學(xué)院院士鄂維南、深勢科技創(chuàng)始人兼首席科學(xué)家和北京科學(xué)智能研究院院長張林峰等。
柯國霖在 ?? 上分享表示:Uni-3DAR 的核心是一種通用的粗到細(xì) token 化方法(coarse-to-fine tokenization),它能將 3D 結(jié)構(gòu)轉(zhuǎn)化為一維的 token 序列。
- 論文標(biāo)題:Uni-3DAR: Unified 3D Generation and Understanding via Autoregression on Compressed Spatial Tokens
- 論文地址:https://arxiv.org/abs/2503.16278
- 項目主頁:https://uni-3dar.github.io
- 代碼倉庫:https://github.com/dptech-corp/Uni-3DAR
基于這套通用的 token 化方法,Uni-3DAR 使用自回歸的方式,統(tǒng)一了 3D 結(jié)構(gòu)的生成和理解任務(wù)。大量實驗表明,Uni-3DAR 在分子生成、晶體結(jié)構(gòu)生成與預(yù)測、蛋白結(jié)合位點預(yù)測、分子對接和分子預(yù)訓(xùn)練等多個任務(wù)中均取得了領(lǐng)先性能。尤其在生成任務(wù)中,相較于現(xiàn)有的擴(kuò)散模型,其性能實現(xiàn)了高達(dá) 256% 的相對提升,推理速度提升達(dá) 21.8 倍,充分驗證了該框架的有效性與高效性。此外,此模型不僅可以用在微觀的 3D 分子,也可以用到宏觀的 3D 任務(wù)上,具備跨尺度的能力。
具體來說,Uni-3DAR 解決了 3D 結(jié)構(gòu)建模里的兩個痛點:
第一,數(shù)據(jù)表示不統(tǒng)一。當(dāng)前的 3D 結(jié)構(gòu)存在多種表示方式,尤其在不同尺度下差異顯著。宏觀結(jié)構(gòu)常用點云、網(wǎng)格(Mesh)等表示方式,而微觀結(jié)構(gòu)則多采用原子坐標(biāo)或圖結(jié)構(gòu)。這些表示方式的差異導(dǎo)致建模思路截然不同。即使在同一尺度,由于數(shù)據(jù)特性的差異,不同類型的結(jié)構(gòu)(如晶體、蛋白質(zhì)、分子)也往往采用各自專用的表示與模型,難以兼容。這種表示上的割裂嚴(yán)重限制了模型的通用性,也阻礙了構(gòu)建可借助大規(guī)模數(shù)據(jù)訓(xùn)練的通用基礎(chǔ)模型的可能性。
第二,建模任務(wù)不統(tǒng)一。 3D 結(jié)構(gòu)相關(guān)任務(wù)可分為生成和理解兩大類,但它們各自獨(dú)立發(fā)展。生成任務(wù)多依賴擴(kuò)散模型,從隨機(jī)噪聲逐步合成穩(wěn)定結(jié)構(gòu),而理解任務(wù)則主要基于無監(jiān)督預(yù)訓(xùn)練方法。相比之下,大型語言模型(LLM)已通過自回歸方式成功實現(xiàn)了生成與理解任務(wù)的統(tǒng)一,但這種統(tǒng)一范式在 3D 結(jié)構(gòu)建模領(lǐng)域仍然鮮有嘗試。若能借助自回歸方法統(tǒng)一 3D 任務(wù)建模,不僅有望打通理解與生成的界限,更可能將 3D 結(jié)構(gòu)納入多模態(tài)大語言模型的處理范式,繼圖像和視頻之后成為 LLM 可理解的新模態(tài),為構(gòu)建面向物理世界的通用多模態(tài)科學(xué)模型奠定基礎(chǔ)。
Uni-3DAR 整體架構(gòu)
下面我們具體解讀下這篇論文的兩個核心技術(shù)。
Compressed Spatial Tokens
統(tǒng)一微觀與宏觀 3D 結(jié)構(gòu)
3D 結(jié)構(gòu)在微觀(如原子、分子、蛋白質(zhì))和宏觀(如物體整體、力學(xué)結(jié)構(gòu))層面均表現(xiàn)出顯著稀疏性:大部分空間為空白,只有局部區(qū)域含有重要信息。傳統(tǒng)的全體素網(wǎng)格表示計算資源消耗巨大,無法利用這種稀疏性。
為此,Uni-3DAR 提出了一種層次化、由粗到細(xì)的 token 化方法,實現(xiàn)了數(shù)據(jù)的高效壓縮和統(tǒng)一表示,既適用于微觀也適用于宏觀 3D 結(jié)構(gòu)建模,為后續(xù)的自回歸生成與理解任務(wù)提供了堅實基礎(chǔ)。
1. 層次化八叉樹壓縮
該方法首先利用八叉樹對整個 3D 空間進(jìn)行無損壓縮。具體做法是從包含整個結(jié)構(gòu)的一個大格子開始,針對非空格子(即包含原子或其他結(jié)構(gòu)信息的區(qū)域),遞歸地將其均分為 8 個等大小的子單元。經(jīng)過多層細(xì)分后,形成一個由粗到細(xì)的層次結(jié)構(gòu),其每一層的 token 不僅記錄了區(qū)域是否為空,還保留了該區(qū)域的空間位置信息(由所在層次及格子中心坐標(biāo)確定),為后續(xù)的自回歸生成提供了明確的空間先驗。
2. 精細(xì)結(jié)構(gòu) token 化
雖然八叉樹可以有效壓縮空白區(qū)域,但它僅提供了粗粒度的空間劃分,無法捕捉到諸如原子類型、精確坐標(biāo)(在微觀結(jié)構(gòu)中)或物體表面細(xì)節(jié)(在宏觀結(jié)構(gòu)中)等重要信息。
為此,該團(tuán)隊在最后層非空區(qū)域內(nèi)進(jìn)一步引入了「3D patch」的概念 —— 類似于圖像領(lǐng)域中的 2D patch 的處理。通過將局部結(jié)構(gòu)細(xì)節(jié)進(jìn)行離散化(例如采用向量量化技術(shù)),將連續(xù)的空間信息轉(zhuǎn)化為離散的 token。
這樣一來,無論是描述微觀尺度下單個原子的信息,還是刻畫宏觀尺度下物體表面的細(xì)節(jié),都能以同一形式進(jìn)行表示。
3. 二級子樹壓縮
由于即使在八叉樹結(jié)構(gòu)下,token 數(shù)量仍可能較多,該方法進(jìn)一步提出了二級子樹壓縮策略。具體來說,將一個父節(jié)點及其 8 個子節(jié)點的信息合并為一個單一的 token(利用父節(jié)點固定狀態(tài)以及子節(jié)點的二值特征,共可組合成 256 種狀態(tài)),從而將 token 總數(shù)約降低 8 倍。這不僅大幅提高了計算效率,也為大規(guī)模 3D 結(jié)構(gòu)的高效建模提供了可能。
綜上,該方法充分利用了 3D 結(jié)構(gòu)固有的稀疏性,通過八叉樹分解、精細(xì) token 化與二級子樹壓縮,不僅大幅降低了數(shù)據(jù)表示的復(fù)雜度,而且實現(xiàn)了從微觀到宏觀 3D 結(jié)構(gòu)的統(tǒng)一表示,為后續(xù)自回歸生成與理解任務(wù)提供了高效、通用的數(shù)據(jù)基礎(chǔ)。
Masked Next-Token Preiction
統(tǒng)一生成和理解的自回歸框架
在傳統(tǒng)自回歸模型中,token 的位置是固定的 —— 例如在文本生成中,第 i 個 token 后總是緊接著第 i+1 個 token,因此下一個 token 的位置可以直接推斷,無需顯式建模。
然而,在該論文提出的粗到細(xì) 3D token 化方法中,token 是動態(tài)展開的,其位置在不同樣本間存在較大變化;如果不顯式提供位置信息,自回歸預(yù)測的難度將大大增加。為此,該論文提出了 Masked Next-Token Prediction 策略。
具體而言,該方法對每個 token 復(fù)制一份,確保兩個副本具有相同的位置信息,然后將其中一個副本替換為 [MASK] token。在自回歸預(yù)測過程中,由于被掩碼 token 與目標(biāo) token 的位置信息完全一致,模型能夠直接利用這一明確的位置信息來預(yù)測下一個 token 的內(nèi)容,從而更精確地捕捉下一個 token 的位置特征,提高預(yù)測效果。盡管復(fù)制 token 使序列長度翻倍,但實驗結(jié)果表明,該策略顯著提升了性能,而推理速度僅下降 15% 至 30%。
基于 Masked Next-Token Prediction,該論文構(gòu)建了一個統(tǒng)一的自回歸框架,使得 3D 結(jié)構(gòu)的生成與理解任務(wù)能夠在單一模型內(nèi)同時進(jìn)行。
具體來說,生成任務(wù)(包括單幀與多幀生成)在被掩碼的 token 上執(zhí)行,利用自回歸機(jī)制逐步構(gòu)建結(jié)構(gòu);token 級理解任務(wù)(如原子級屬性預(yù)測)則依托精細(xì)結(jié)構(gòu) token 進(jìn)行;而結(jié)構(gòu)級理解任務(wù)則引入了一個特殊的 [EoS](End of Structure) token,用于捕捉整體結(jié)構(gòu)的全局信息。
此設(shè)計使不同任務(wù)對應(yīng)的 token 在模型內(nèi)部彼此獨(dú)立、互不干擾,從而支持聯(lián)合訓(xùn)練。同時,自回歸特性也便于將其他模態(tài)數(shù)據(jù)(例如自然語言文本、蛋白質(zhì)序列、儀器信號等)統(tǒng)一到單個模型,進(jìn)一步提升模型的泛化能力和實用性。
實驗結(jié)果
該論文在微觀 3D 結(jié)構(gòu)領(lǐng)域設(shè)計了一系列任務(wù),包括分子生成、晶體結(jié)構(gòu)生成與預(yù)測、蛋白結(jié)合位點預(yù)測、蛋白小分子對接以及基于預(yù)訓(xùn)練的分子性質(zhì)預(yù)測。
實驗結(jié)果顯示,在生成任務(wù)中,Uni-3DAR 的性能大幅超過了擴(kuò)散模型方法;而在無監(jiān)督預(yù)訓(xùn)練的理解任務(wù)上,其表現(xiàn)與基于雙向注意力的模型基本持平。這些成果充分證明,Uni-3DAR 不僅能統(tǒng)一不同類型的 3D 結(jié)構(gòu)數(shù)據(jù)及任務(wù),而且在效果和速度上均實現(xiàn)了顯著提升。
3D 小分子生成任務(wù)性能
晶體結(jié)構(gòu)預(yù)測,以及基于多模態(tài)信息(粉末 X 射線衍射譜)的晶體結(jié)構(gòu)解析性能
蛋白結(jié)合位點預(yù)測效果
蛋白小分子對接效果
基于預(yù)訓(xùn)練的小分子屬性預(yù)測效果,其中 Uni-Mol 和 SpaceFormer 也為深勢科技提出的專用模型,Uni-3DAR 超過了 Uni-Mol,與 SpaceFormer 基本持平
基于預(yù)訓(xùn)練的高分子聚合物性質(zhì)預(yù)測,其中 Uni-Mol 和 MMPolymer 也為深勢科技提出的專用模型,Uni-3DAR 超過了 Uni-Mol,與 MMPolymer 基本持平
未來展望
目前,Uni-3DAR 的實驗主要集中在微觀結(jié)構(gòu)領(lǐng)域,因此亟需在宏觀 3D 結(jié)構(gòu)任務(wù)中進(jìn)一步驗證其通用性和擴(kuò)展性。
此外,為保證與以往工作的公平對比,當(dāng)前 Uni-3DAR 在每個任務(wù)上均采用獨(dú)立訓(xùn)練。未來的一個重要方向是融合多種數(shù)據(jù)類型與任務(wù),構(gòu)建并聯(lián)合訓(xùn)練一個更大規(guī)模的 Uni-3DAR 基座模型,以進(jìn)一步提升性能與泛化能力。
同時,Uni-3DAR 還具備天然的多模態(tài)擴(kuò)展?jié)摿?。后續(xù)可以引入更多模態(tài)的信息,例如蛋白質(zhì)序列、氨基酸組成,甚至結(jié)合大語言模型與科學(xué)文獻(xiàn)知識,共同訓(xùn)練一個具備物理世界理解能力的多模態(tài)科學(xué)語言模型,從而為構(gòu)建通用科學(xué)智能體打下基礎(chǔ)。