AI4Science的基石：幾何圖神經網絡，最完整綜述來了！人大高瓴聯合騰訊AI lab、清華、史丹佛等發布

編輯 | XS

Nature 在 2023 年 11 月發表了兩項重要研究成果：蛋白質合成技術 Chroma 和晶體材料設計方法 GNoME。這兩項研究都採用了圖神經網路作為處理科學資料的工具。

實際上，圖神經網絡，特別是幾何圖神經網絡，一直是科學智慧（AI for Science）研究的重要工具。這是因為，科學領域中的粒子、分子、蛋白質、晶體等物理系統都可以被建模成一種特殊的資料結構──幾何圖。

與一般的拓樸圖不同，為了更好地描述物理系統，幾何圖加入了不可或缺的空間信息，需要滿足平移、旋轉和翻轉的物理對稱性。鑑於幾何圖神經網路對於物理系統建模的優越性，近年來各類方法層出不窮，論文數量持續成長。

近日，人大高瓴聯合騰訊 AI Lab、清華、史丹佛等機構發布綜述論文：《A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications》。此綜述在簡要介紹群論、對稱性等理論知識的基礎上，從資料結構、模型到眾多科學應用，對相關幾何圖神經網路文獻進行了系統性的梳理。

論文連結：https://arxiv.org/abs/2403.00485

GitHub連結：https:/ /github.com/RUC-GLAD/GGNN4Science

在這篇綜述中，作者研究了300 多篇參考文獻，歸納出3 種不同的幾何圖神經網路模型，介紹了面向粒子、分子、蛋白質等多種科學數據上共23 種不同任務的相關方法，收集了50 多個相關評測資料集。最後，綜述展望了未來的研究方向，包括幾何圖基礎模型、與大語言模型結合等。

以下是各章節簡單介紹。

幾何圖資料結構

幾何圖由鄰接矩陣、節點特徵、節點幾何資訊（例如座標）構成。在歐氏空間中，幾何圖通常表現出平移、旋轉和反射的物理對稱性，一般使用群來刻畫這些變換，包括歐式群、平移群、正交群、置換群等等。直觀上看，可以理解為置換、平移、旋轉、翻轉四種操作依一定順序的​​複合。

對於眾多 AI for Science 領域，幾何圖是一種有力且通用的表示方法，其可用於表示眾多物理系統，包括小分子、蛋白質、晶體、物理點雲等。

幾何圖神經網路模型

#根據實際問題中的求解目標對於對稱性的要求，本文將幾何圖神經網路分為三類：不變（invariant）模型、等變（equivariant）模型、以及受Transformer 架構啟發的Geometric Graph Transformer，其中等變模型又細分為標量化方法模型（Scalarization-Based Model）與基於球面調和的高階可操控模型（High-Degree Steerable Model）。按照上述規則，文章收集並歸類了近年來知名的幾何圖神經網路模型。

這裡我們透過各分支的代表性工作簡單介紹不變模型（SchNet[1]）、標量化方法模型（EGNN[2]）、高階可操控模型（TFN[3]）的關聯與差異。可以發現三者皆是採用了訊息傳遞機制，只是身為等變模型的後兩者額外引入了一次幾何訊息傳遞。

不變模型主要利用節點本身的特徵（如原子種類、質量、帶電量等）與原子間不變的特徵（如距離、角度[4]、二面角[5]）等進行訊息計算，隨後進行傳播。

而在此之上，標量化方法額外透過節點間座標差引入了幾何訊息，並將不變資訊作為幾何資訊的權重進行線性組合，實現了等變性的引入。

高階可操控模型則是使用了高階的球面調和（Spherical Harmonics）與Wigner-D 矩陣表徵系統的幾何信息，這類方法透過量子力學中的Clebsch–Gordan 係數操控不可約表示的階數，從而實現幾何訊息傳遞過程。

幾何圖神經網路透過這類設計保證的對稱性，準確率有大幅提升，並且在生成任務中也大放異彩。

下圖是幾何圖神經網路與傳統模型在QM9、PDBBind、SabDab 三個資料集上進行分子性質預測、蛋白質-配體對接和抗體設計（生成）三個任務中的結果，可以明顯看出幾何圖神經網路的優勢。

科學應用

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

上表展示了各個領域的常見任務與經典模型，其中，按照單一實例與多實例（如化學反應，需要多分子共同參與），文章單獨區分了小分子-小分子、小分子-蛋白質、蛋白質-蛋白質三個領域。

為了更好地方便領域內進行模型設計與實驗開展，文章依照單一實例與多實例統計了兩類任務的常用資料集與基準（benchmark），並記錄了不同資料集的樣本量與任務種類。

下表整理了常見的單一實例任務資料集。

下表整理了常見的多實例任務資料集。

未來展望

文章就幾個面向進行了初步的展望，希望能作拋磚引玉之用：

1. 幾何圖基礎模型

在各種任務和領域中採用統一的基礎模型的優越性在GPT系列模型的顯著進步中已經體現得淋漓盡致。如何在任務空間、資料空間、模型空間進行合理的設計，從而將這種想法引入到針對幾何圖神經網路的設計上仍是一個有趣的開放問題。

2. 模型訓練與現實世界實驗驗證的高效循環

科學資料的獲取是昂貴且耗時的，而僅在獨立資料集上評估的模型不能直接反應來自現實世界的回饋。如何類似於GNoME（整合了一個端到端的流水線，包括圖網絡訓練、密度泛函理論計算和用於材料發現和合成的自動實驗室）實現高效的模型-現實循環迭代的實驗範式的重要性將會與日俱增。

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

大型語言模型（LLMs）已被廣泛證明具有豐富的知識，涵蓋了各個領域。雖然已經有一些工作利用 LLMs 進行某些任務，例如分子屬性預測和藥物設計，但它們僅在基元或分子圖上操作。如何將它們與幾何圖神經網路有機組合，使其能夠處理 3D 結構資訊並在 3D 結構上執行預測或生成，仍然具有相當的挑戰性。

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

毫無疑問，等變性對增強資料效率和模型泛化能力至關重要，但值得注意的是，過強等變性約束有時可能過於限制模型，潛在地損害其性能。因此，如何使得所設計的模型在等變性與適應能力中取得平衡是一個非常有趣的問題。這方面的探索不僅可以豐富我們對模型行為的理解，還可以為開發更具穩健性和通用性的解決方案鋪平道路，使其具有更廣泛的適用性。

參考文獻

[1] Schütt K, Kindermans P J, Sauceda Felix H E, et al. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions[ J]. Advances in neural information processing systems, 2017, 30.

[2] Satorras VG, Hoogeboom E, Welling M. E(n)等變圖神經網路[C]//機器學習國際會議。 PMLR，2021：9323-9332。

[3] Thomas N、Smidt T、Kearnes S 等人。張量場網路：3d 點雲的旋轉和平移等變神經網路[J]. arXiv 預印本 arXiv:1802.08219, 2018.

[4] Gasteiger J, Groß J, Günnemann S. 分子圖的定向訊息傳遞[C]//學習表示國際會議。 2019.

[5] Gasteiger J, Becker F, Günnemann S. Gemnet: 分子的通用有向圖神經網路[J].神經資訊處理系統進展，2021, 34: 6790-6802。

[6] Merchant A、Batzner S、Schoenholz S S 等。擴展深度學習以促進材料發現[J].自然, 2023, 624(7990): 80-85.

以上是AI4Science的基石：幾何圖神經網絡，最完整綜述來了！人大高瓴聯合騰訊AI lab、清華、史丹佛等發布的詳細內容。更多資訊請關注PHP中文網其他相關文章！