編輯 | XS
2023 年 11 月,Nature 連續刊登了兩篇重大成果:蛋白質生成方法 Chroma 和晶體材料設計方法 GNoME,均使用了圖神經網絡作為科學數據的表示工具。
實際上,圖神經網絡,特別是幾何圖神經網絡,一直是科學智能(AI for Science)研究的重要工具。這是因為,科學領域中的粒子、分子、蛋白質、晶體等物理系統均可被建模成一種特殊的數據結構——幾何圖。
與一般的拓撲圖不同,為了更好描述物理系統,幾何圖加入了不可或缺的空間信息,需要滿足平移、旋轉和翻轉的物理對稱性。鑒于幾何圖神經網絡對于物理系統建模的優越性,近年來各類方法層出不窮,論文數量持續增長。
近日,人大高瓴聯合騰訊 AI Lab、清華、斯坦福等機構發布綜述論文:《A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications》。該綜述在簡要介紹群論、對稱性等理論知識的基礎上,從數據結構、模型到眾多科學應用,對相關幾何圖神經網絡文獻進行了系統的梳理。
在這篇綜述中,作者調研了 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. 等變性約束條件的放松
毫無疑問,等變性對增強數據效率和模型泛化能力至關重要,但值得注意的是,過強等變性約束有時可能過于限制模型,潛在地損害其性能。因此,如何使得所設計的模型在等變性與適應能力中取得平衡是一個非常有趣的問題。這方面的探索不僅可以豐富我們對模型行為的理解,還可以為開發更具魯棒性和通用性的解決方案鋪平道路,使其具有更廣泛的適用性。
參考文獻
原創文章,作者:計算搬磚工程師,如若轉載,請注明來源華算科技,注明出處:http://www.zzhhcy.com/index.php/2024/03/08/0368edb544/