編輯 | XS

2023 年 11 月，Nature 連續(xù)刊登了兩篇重大成果：蛋白質(zhì)生成方法 Chroma 和晶體材料設(shè)計(jì)方法 GNoME，均使用了圖神經(jīng)網(wǎng)絡(luò)作為科學(xué)數(shù)據(jù)的表示工具。

實(shí)際上，圖神經(jīng)網(wǎng)絡(luò)，特別是幾何圖神經(jīng)網(wǎng)絡(luò)，一直是科學(xué)智能（AI for Science）研究的重要工具。這是因?yàn)椋茖W(xué)領(lǐng)域中的粒子、分子、蛋白質(zhì)、晶體等物理系統(tǒng)均可被建模成一種特殊的數(shù)據(jù)結(jié)構(gòu)——幾何圖。

與一般的拓?fù)鋱D不同，為了更好描述物理系統(tǒng)，幾何圖加入了不可或缺的空間信息，需要滿足平移、旋轉(zhuǎn)和翻轉(zhuǎn)的物理對稱性。鑒于幾何圖神經(jīng)網(wǎng)絡(luò)對于物理系統(tǒng)建模的優(yōu)越性，近年來各類方法層出不窮，論文數(shù)量持續(xù)增長。

近日，人大高瓴聯(lián)合騰訊 AI Lab、清華、斯坦福等機(jī)構(gòu)發(fā)布綜述論文：《A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications》。該綜述在簡要介紹群論、對稱性等理論知識的基礎(chǔ)上，從數(shù)據(jù)結(jié)構(gòu)、模型到眾多科學(xué)應(yīng)用，對相關(guān)幾何圖神經(jīng)網(wǎng)絡(luò)文獻(xiàn)進(jìn)行了系統(tǒng)的梳理。

在這篇綜述中，作者調(diào)研了 300 多篇參考文獻(xiàn)，歸納出 3 種不同的幾何圖神經(jīng)網(wǎng)絡(luò)模型，介紹了面向粒子、分子、蛋白質(zhì)等多種科學(xué)數(shù)據(jù)上共 23 種不同任務(wù)的相關(guān)方法，收集了 50 多個(gè)相關(guān)評測數(shù)據(jù)集。最后，綜述展望了未來的研究方向，包括幾何圖基礎(chǔ)模型、與大語言模型結(jié)合等。

下面是各章節(jié)簡要介紹。

幾何圖數(shù)據(jù)結(jié)構(gòu)

幾何圖由鄰接矩陣、節(jié)點(diǎn)特征、節(jié)點(diǎn)幾何信息（例如坐標(biāo)）構(gòu)成。在歐氏空間中，幾何圖通常表現(xiàn)出平移、旋轉(zhuǎn)和反射的物理對稱性，一般使用群來刻畫這些變換，包括歐式群、平移群、正交群、置換群等等。直觀上看，可以理解為置換、平移、旋轉(zhuǎn)、翻轉(zhuǎn)四種操作按一定順序的復(fù)合。

對于眾多 AI for Science 領(lǐng)域，幾何圖是一種有力且通用的表示方法，其可以用于表示眾多物理系統(tǒng)，包括小分子、蛋白質(zhì)、晶體、物理點(diǎn)云等。

幾何圖神經(jīng)網(wǎng)絡(luò)模型

根據(jù)實(shí)際問題中的求解目標(biāo)對于對稱性的要求，本文將幾何圖神經(jīng)網(wǎng)絡(luò)分為三類：不變（invariant）模型、等變（equivariant）模型、以及受 Transformer 架構(gòu)啟發(fā)的 Geometric Graph Transformer，其中等變模型又細(xì)分為標(biāo)量化方法模型（Scalarization-Based Model）與基于球面調(diào)和的高階可操控模型（High-Degree Steerable Model）。按照上述規(guī)則，文章收集并歸類了近年來知名的幾何圖神經(jīng)網(wǎng)絡(luò)模型。

這里我們通過各個(gè)分支的代表性工作簡要介紹不變模型（SchNet[1]）、標(biāo)量化方法模型（EGNN[2]）、高階可操控模型（TFN[3]）的關(guān)聯(lián)與區(qū)別?？梢园l(fā)現(xiàn)三者均是采用了消息傳遞機(jī)制，只是身為等變模型的后兩者額外引入了一次幾何消息傳遞。

不變模型主要利用節(jié)點(diǎn)本身的特征（如原子種類、質(zhì)量、帶電量等）與原子間的不變特征（如距離、角度[4]、二面角[5]）等進(jìn)行消息計(jì)算，隨后進(jìn)行傳播。

而在此之上，標(biāo)量化方法額外通過節(jié)點(diǎn)間坐標(biāo)差引入了幾何信息，并將不變信息作為幾何信息的權(quán)重進(jìn)行線性組合，實(shí)現(xiàn)了等變性的引入。

高階可操控模型則是使用了高階的球面調(diào)和（Spherical Harmonics）與 Wigner-D 矩陣表征系統(tǒng)的幾何信息，這類方法通過量子力學(xué)中的 Clebsch–Gordan 系數(shù)操控不可約表示的階數(shù)，從而實(shí)現(xiàn)幾何消息傳遞過程。

幾何圖神經(jīng)網(wǎng)絡(luò)通過這類設(shè)計(jì)保證的對稱性，準(zhǔn)確率有大幅提升，并且在生成任務(wù)中也大放異彩。

下圖是幾何圖神經(jīng)網(wǎng)絡(luò)與傳統(tǒng)模型在 QM9、PDBBind、SabDab 三個(gè)數(shù)據(jù)集上進(jìn)行分子性質(zhì)預(yù)測、蛋白質(zhì)-配體對接和抗體設(shè)計(jì)（生成）三個(gè)任務(wù)中的結(jié)果，可以明顯看出幾何圖神經(jīng)網(wǎng)絡(luò)的優(yōu)勢。

科學(xué)應(yīng)用

在科學(xué)應(yīng)用方面，綜述涵蓋了物理（粒子）、生物化學(xué)（小分子、蛋白質(zhì)）以及其它如晶體等多個(gè)應(yīng)用場景，任務(wù)定義與所需保證對稱性種類出發(fā)，分別介紹了各個(gè)任務(wù)中的常用數(shù)據(jù)集與該類任務(wù)中的經(jīng)典模型設(shè)計(jì)思路。

上表展示了各個(gè)領(lǐng)域的常見任務(wù)與經(jīng)典模型，其中，按照單一實(shí)例與多實(shí)例（如化學(xué)反應(yīng)，需要多分子共同參與），文章單獨(dú)區(qū)分了小分子-小分子、小分子-蛋白質(zhì)、蛋白質(zhì)-蛋白質(zhì)三個(gè)領(lǐng)域。

為了更好地方便領(lǐng)域內(nèi)進(jìn)行模型設(shè)計(jì)與實(shí)驗(yàn)開展，文章按照單一實(shí)例與多實(shí)例統(tǒng)計(jì)了兩類任務(wù)的常用數(shù)據(jù)集與基準(zhǔn)（benchmark），并記錄了不同數(shù)據(jù)集的樣本量與任務(wù)種類。

下表整理了常見的單實(shí)例任務(wù)數(shù)據(jù)集。

下表整理了常見的多實(shí)例任務(wù)數(shù)據(jù)集。

未來展望

文章就幾個(gè)方面進(jìn)行了初步的展望，希望能作拋磚引玉之用：

1. 幾何圖基礎(chǔ)模型

在各種任務(wù)和領(lǐng)域中采用統(tǒng)一的基礎(chǔ)模型的優(yōu)越性在GPT系列模型的顯著進(jìn)步中已經(jīng)體現(xiàn)得淋漓盡致。如何在任務(wù)空間、數(shù)據(jù)空間、模型空間進(jìn)行合理的設(shè)計(jì)，從而將這種思路引入到針對幾何圖神經(jīng)網(wǎng)絡(luò)的設(shè)計(jì)上仍是一個(gè)有趣的開放問題。

2. 模型訓(xùn)練與現(xiàn)實(shí)世界實(shí)驗(yàn)驗(yàn)證的高效循環(huán)

科學(xué)數(shù)據(jù)的獲取是昂貴且耗時(shí)的，而僅僅在獨(dú)立數(shù)據(jù)集上評估的模型不能直接反應(yīng)來自現(xiàn)實(shí)世界的反饋。如何類似于GNoME（集成了一個(gè)端到端的流水線，包括圖網(wǎng)絡(luò)訓(xùn)練、密度泛函理論計(jì)算和用于材料發(fā)現(xiàn)和合成的自動實(shí)驗(yàn)室）實(shí)現(xiàn)高效的模型-現(xiàn)實(shí)循環(huán)迭代的實(shí)驗(yàn)范式的重要性將會與日俱增。

3. 與大型語言模型（LLMs）的融合

大型語言模型（LLMs）已被廣泛證明具有豐富的知識，涵蓋了各個(gè)領(lǐng)域。雖然已經(jīng)有一些工作利用 LLMs 進(jìn)行某些任務(wù)，例如分子屬性預(yù)測和藥物設(shè)計(jì)，但它們僅在基元或分子圖上操作。如何將它們與幾何圖神經(jīng)網(wǎng)絡(luò)有機(jī)組合，使其能夠處理 3D 結(jié)構(gòu)信息并在 3D 結(jié)構(gòu)上執(zhí)行預(yù)測或生成，仍然具有相當(dāng)?shù)奶魬?zhàn)性。

4. 等變性約束條件的放松

毫無疑問，等變性對增強(qiáng)數(shù)據(jù)效率和模型泛化能力至關(guān)重要，但值得注意的是，過強(qiáng)等變性約束有時(shí)可能過于限制模型，潛在地?fù)p害其性能。因此，如何使得所設(shè)計(jì)的模型在等變性與適應(yīng)能力中取得平衡是一個(gè)非常有趣的問題。這方面的探索不僅可以豐富我們對模型行為的理解，還可以為開發(fā)更具魯棒性和通用性的解決方案鋪平道路，使其具有更廣泛的適用性。

參考文獻(xiàn)