編輯 | XS
2023 年 11 月，Nature 連續刊登了兩篇重大成果：蛋白質生成方法 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 V G, Hoogeboom E, Welling M. E (n) equivariant graph neural networks[C]//International conference on machine learning. PMLR, 2021: 9323-9332.
[3] Thomas N, Smidt T, Kearnes S, et al. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds[J]. arXiv preprint arXiv:1802.08219, 2018.
[4] Gasteiger J, Groß J, Günnemann S. Directional Message Passing for Molecular Graphs[C]//International Conference on Learning Representations. 2019.
[5] Gasteiger J, Becker F, Günnemann S. Gemnet: Universal directional graph neural networks for molecules[J]. Advances in Neural Information Processing Systems, 2021, 34: 6790-6802.
[6] Merchant A, Batzner S, Schoenholz S S, et al. Scaling deep learning for materials discovery[J]. Nature, 2023, 624(7990): 80-85.