GNN for Science: 騰訊AI Lab、清華共同發(fā)文綜述等變圖神經(jīng)網(wǎng)絡(luò)
近年來(lái),越來(lái)越多的人工智能方法在解決傳統(tǒng)自然科學(xué)等問(wèn)題上大放異彩, 在一些重要的學(xué)科問(wèn)題(例如蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè))上取得了令人矚目的進(jìn)展。在物理領(lǐng)域的研究中,非常多的物理問(wèn)題都會(huì)涉及建模物體的的一些幾何特征,例如空間位置,速度,加速度等。這種特征往往可以使用幾何圖這一形式來(lái)表示。不同于一般的圖數(shù)據(jù),幾何圖一個(gè)非常重要的特征是額外包含旋轉(zhuǎn),平移,翻轉(zhuǎn)對(duì)稱性。這些對(duì)稱性往往反應(yīng)了某些物理問(wèn)題的本質(zhì)。因此,最近以來(lái),大量工作利用了幾何圖中的對(duì)稱性,基于經(jīng)典圖神網(wǎng)絡(luò)設(shè)計(jì)了很多具有等變性質(zhì)的模型去解決對(duì)幾何圖建模問(wèn)題。盡管在這一領(lǐng)域,等變圖神經(jīng)網(wǎng)絡(luò)模型取得了長(zhǎng)足的發(fā)展,但是還缺乏一個(gè)系統(tǒng)性的對(duì)這一領(lǐng)域的調(diào)研。為此,騰訊 AI Lab, 清華 AIR & 計(jì)算機(jī)系在綜述:《Geometrically Equivariant Graph Neural Networks: A Survey》中,對(duì)等變圖神經(jīng)網(wǎng)絡(luò)的結(jié)構(gòu)和相關(guān)任務(wù)進(jìn)行了一個(gè)系統(tǒng)梳理。
綜述論文鏈接:https://arxiv.org/abs/2202.07230
在物理和化學(xué)領(lǐng)域,很多問(wèn)題需要去處理帶有幾何特征的圖。例如,化學(xué)小分子和蛋白質(zhì)都可以建模成一個(gè)有原子和其化學(xué)鍵關(guān)系組成的幾何圖。在這個(gè)圖中,除了包含原子的一些內(nèi)在特征以外,我們還需要考慮到每個(gè)原子在空間的三維坐標(biāo)這一幾何特征。而在物理學(xué)的多體問(wèn)題中,每個(gè)粒子的幾何特征則包括坐標(biāo),速度,旋轉(zhuǎn)等。不同于一般特征,這些幾何特征往往都具備著一些對(duì)稱性和等變性。正因?yàn)槿绱耍趯?duì)對(duì)稱性的建模,大量基于圖神經(jīng)網(wǎng)絡(luò)的改進(jìn)模型在近年來(lái)被提出。這一類模型,因?yàn)榭朔藗鹘y(tǒng)圖神經(jīng)網(wǎng)絡(luò)無(wú)法很好處理這類具有等變對(duì)稱性質(zhì)的特征的缺點(diǎn),被統(tǒng)稱為等變圖神經(jīng)網(wǎng)絡(luò)。
在這篇綜述里面,我們系統(tǒng)性的梳理了近年等變圖神經(jīng)網(wǎng)絡(luò)的發(fā)展脈絡(luò),并且提供了一個(gè)簡(jiǎn)潔的視角幫助讀者能夠很快的理解這類網(wǎng)絡(luò)的內(nèi)涵?;谙鞑ズ途酆虾瘮?shù)的不同,我們將現(xiàn)有的等變圖神經(jīng)網(wǎng)絡(luò)分為三類。與此同時(shí),我們還詳盡闡釋了當(dāng)前的挑戰(zhàn)和未來(lái)的可能方向。
等變圖神經(jīng)網(wǎng)絡(luò)基礎(chǔ)框架在實(shí)際應(yīng)用中,我們需要處理的圖不僅包含拓?fù)溥B接和節(jié)點(diǎn)特征,同時(shí)也會(huì)包含一些幾何特征。在使用圖神經(jīng)網(wǎng)絡(luò)處理這些數(shù)據(jù)的時(shí)候,不同的特征需要滿足不同的性質(zhì)。例如,在預(yù)測(cè)分子的能量時(shí),我們需要這個(gè)預(yù)測(cè)對(duì)于輸入的幾何特征是不變的,而在分子動(dòng)力學(xué)應(yīng)用中,我們則需要預(yù)測(cè)的結(jié)果和輸入的幾何特征是等變的。為了達(dá)到這樣的目的,我們提出了一個(gè)等變圖神經(jīng)網(wǎng)絡(luò)的通用框架:
在這個(gè)框架中,代表輸入圖的幾何特征,而h_i,h_j 代表非幾何特征。
和
分別代表在邊(x,j)上的幾何和非幾何的消息。
和
分別是針對(duì)幾何和非幾何消息的聚合函數(shù)。除此以外, 針對(duì)非幾何信息的消息函數(shù)
是對(duì)于輸入
來(lái)說(shuō) G - 不變的。而針對(duì)幾何信息的
消息函數(shù)則是對(duì)于輸入
來(lái)說(shuō) G - 等變的。下圖展示了這一通用框架的操作:
基于以上的通用框架,我們?cè)谙卤砜偨Y(jié)了當(dāng)前主流的等變圖神經(jīng)網(wǎng)絡(luò)模型。與此同時(shí),基于消息表示的類別不同,我們將現(xiàn)有的等變圖神經(jīng)網(wǎng)絡(luò)模型分為三大類:不可約表示(Irreducible Representation),正則表示(Regular Representation) 和標(biāo)量化(Scalarization ).
基于不可約表示信息的模型這類模型基于表示論中關(guān)于緊群的線性表示可以拆解為一系列的不可約表示的直積這一理論。從而在 SE(3) 群中構(gòu)建滿足等變性質(zhì)的消息模型。例如,在 TFN 中:
TFN 層利用 Clebsch-Gordan 系數(shù)的性質(zhì)構(gòu)造了一個(gè)對(duì)于任意旋轉(zhuǎn)參數(shù),對(duì)于任意屬于 SO(3)的旋轉(zhuǎn)操作都等變。有大量的工作基于 TFN 結(jié)構(gòu)做了相應(yīng)的擴(kuò)展,例如加入 Attention 機(jī)制, 引入非線性的 Clebsch-Gordan 系數(shù)等。但是這類方法計(jì)算復(fù)雜度都較高,且不可約表示僅僅適用于特定的群。這約束了這類模型的表達(dá)能力。
基于正則表示信息的模型另一類的工作嘗試?yán)萌旱恼齽t表示來(lái)構(gòu)造群卷積操作。這里的代表性工作李卷積(LieConv)通過(guò) Lifting 操作將輸入映射到群中的元素,然后利用 PointConv 完成群卷積的離散化計(jì)算。在我們的符號(hào)約定下,李卷積可以表示為:
其中是
在群中的映射元素,log 將群元映射到對(duì)應(yīng)的李代數(shù),
是一個(gè) MLP。通過(guò)這種構(gòu)造,李卷積中對(duì)于h_i的更新實(shí)現(xiàn)了對(duì)于任意李群以及其離散子群的不變性。LieTransformer 基于此思想,引入了自注意力機(jī)制來(lái)進(jìn)一步提高模型的性能。基于李群正則表示的模型在群的選取上更加靈活,但是由于要進(jìn)行離散化和采樣,需要在效率和性能之間做出權(quán)衡。同時(shí),以上的更新只考慮了標(biāo)量信息 h,但難以直接推廣到對(duì)幾何信息 x 的更新,除非綜合哈密頓網(wǎng)絡(luò)等工作中的更新方法。
標(biāo)量化去基于群表示論的途徑外,很多工作采用了一種基于標(biāo)量化的建模等變性質(zhì)的方法。這類標(biāo)量化的方法先將幾何特征轉(zhuǎn)化為一些不變的標(biāo)量,然后利用 MLP 等網(wǎng)絡(luò)結(jié)構(gòu)來(lái)得到一個(gè)標(biāo)量變化,最后將這個(gè)變化加回到原有的幾何特征上從而得到等變性。這種標(biāo)量化的方法最早是由 SchNet 和 DimNet 提出,不過(guò)僅僅考慮了模型不變的部分。SphereNet 在之前工作的基礎(chǔ)上進(jìn)一步考慮了在消息傳播網(wǎng)絡(luò)上扭轉(zhuǎn)角的變化。EGNN 作為在標(biāo)量化里面一個(gè)重要的工作,提出了一個(gè)非常靈活的框架:
其中,是對(duì)幾何特征的標(biāo)量化,
函數(shù)為不同的 MLP,通過(guò)將幾何信息和非幾何信息消息進(jìn)行關(guān)聯(lián)
, EGNN 可以同時(shí)保證非幾何特征和幾何特征傳播過(guò)程中的等變性。這個(gè)構(gòu)造結(jié)合了物理知識(shí),可以看成是對(duì)兩個(gè)粒子的庫(kù)倫力 / 重力的計(jì)算的建模。在 EGNN 基礎(chǔ)上,GMN 擴(kuò)展了模型可以描述的幾何特征維度,在建模坐標(biāo)信息的同時(shí)也可以同時(shí)引入更多的幾何信息(如速度、加速度、角速度等)并保證等變性。GemNet 則在 DimeNet 基礎(chǔ)上通過(guò)這一通用的表示將一些更豐富的幾何特征,例如二面角等,結(jié)合到消息傳播的過(guò)程中。此外,還存在一類標(biāo)量化的方法,其基于不變的標(biāo)量和等變的向量的乘積仍然是等變的向量這一觀察來(lái)構(gòu)造等變的消息傳播。例如, PaiNN 和 Equivariant Transformer 在不變的 SchNet 上通過(guò)徑向基函數(shù)建模原子的距離將等變的性質(zhì)擴(kuò)展到 SchNet 上。
以下是對(duì)等變圖神經(jīng)網(wǎng)絡(luò)模型的一個(gè)總結(jié)梳理:
等變圖神經(jīng)網(wǎng)絡(luò)應(yīng)用因?yàn)榭梢愿玫慕缀涡畔ⅲ茸儓D神經(jīng)網(wǎng)絡(luò)在從物理系統(tǒng)到化學(xué)物質(zhì)的各種類型的現(xiàn)實(shí)世界幾何數(shù)據(jù)中具有廣泛的應(yīng)用。這篇綜述簡(jiǎn)單介紹其在物理系統(tǒng),分子數(shù)據(jù)和點(diǎn)云數(shù)據(jù)上的應(yīng)用。下表總結(jié)了現(xiàn)有的等變神經(jīng)網(wǎng)絡(luò)的應(yīng)用方向和數(shù)據(jù)集:
對(duì)復(fù)雜物理系統(tǒng)的建模長(zhǎng)期以來(lái),對(duì)復(fù)雜物理系統(tǒng)的動(dòng)力學(xué)進(jìn)行建模一直是一個(gè)具有挑戰(zhàn)性的話題。在物理系統(tǒng)中,有像帶電粒子這樣的物體,它們通過(guò)基于物理定律的力進(jìn)行相互作用產(chǎn)生運(yùn)動(dòng)軌跡。在 NRI 這篇工作中最早引入 n-body 模擬的問(wèn)題。n-body 中系統(tǒng)包含多個(gè)帶電粒子,這些粒子由相互之間的庫(kù)倫力驅(qū)動(dòng)。n-body 問(wèn)題的目標(biāo)是給定系統(tǒng)的初始條件(坐標(biāo),速度和電荷量)的情況下預(yù)測(cè)這些粒子的動(dòng)力學(xué)軌跡。這一任務(wù)是 E(3)等變的。SE(3)-Transformer 和 EGNN 都展示了等變圖神經(jīng)網(wǎng)絡(luò)在這個(gè)任務(wù)上的潛力。GMN 中進(jìn)一步提出了一個(gè)更有挑戰(zhàn)性的問(wèn)題 -- 帶約束的 n-body 問(wèn)題,即如何在粒子之間有約束,例如連桿或者鉸鏈的情況下對(duì)粒子運(yùn)動(dòng)軌跡做出有效的預(yù)測(cè)。除了微觀方向的數(shù)據(jù)外,NRI 和 GMN 也采用了人體運(yùn)動(dòng)捕捉的宏觀數(shù)據(jù)驗(yàn)證了模型的有效性。
對(duì)分子的建模等變圖神經(jīng)網(wǎng)絡(luò)另一個(gè)重要的應(yīng)用方向是對(duì)分子數(shù)據(jù)的建模。在分子數(shù)據(jù)中,原子的相互作用是有一系列復(fù)雜的化學(xué)物理機(jī)制所決定的。對(duì)于典型的分子數(shù)據(jù),原子非幾何特征往往包含原子本身的一些特征,而幾何特征則是原子的空間坐標(biāo),速度等。原子之間的邊則由化學(xué)鍵或者根據(jù)實(shí)際距離做截?cái)鄟?lái)構(gòu)造。經(jīng)典的在分子上的應(yīng)用包括,分子預(yù)測(cè)和分子生成。 分子預(yù)測(cè):具體來(lái)說(shuō),分子預(yù)測(cè)包含對(duì)分子的屬性和結(jié)構(gòu)的一些預(yù)測(cè)任務(wù)。在分子預(yù)測(cè)這一領(lǐng)域,包含以下經(jīng)典數(shù)據(jù)集。在小分子方面:QM9 是一個(gè)包含 12 個(gè)量子特征預(yù)測(cè)任務(wù)的經(jīng)典小分子數(shù)據(jù)集。M17 則是一個(gè)在 8 個(gè)小分子上得到的動(dòng)態(tài)軌跡的數(shù)據(jù)集,其中還包含了對(duì)應(yīng)狀態(tài)的能量和相互力作用等信息。ISO17 則是一個(gè)類似的分子動(dòng)態(tài)軌跡數(shù)據(jù)集,其包含了 129 個(gè)同分異構(gòu)體的軌跡信息。The Open Catalyst 2020 (OC20) 則包含了催化劑和底物的催化過(guò)程的狀態(tài)信息,其目標(biāo)是給定初始狀態(tài)預(yù)測(cè)目標(biāo)結(jié)構(gòu)和對(duì)應(yīng)狀態(tài)的能量。在大分子方面:MDAnalysis 是一個(gè)較為完備的關(guān)于蛋白質(zhì)級(jí)別的分子動(dòng)力學(xué)模擬的數(shù)據(jù)。Atom3D 是一個(gè)綜合數(shù)據(jù)集,包含 8 個(gè)具有幾何信息的分子預(yù)測(cè)任務(wù),范圍從小分子到 RNA 和蛋白質(zhì)。分子生成:在分子生成領(lǐng)域,等變圖神經(jīng)網(wǎng)絡(luò)往往用于和分子構(gòu)象相關(guān)的生成。ConfGF 和 DSGM 基于旋轉(zhuǎn) - 平移不變 GNN 來(lái)參數(shù)化打分函數(shù)并且構(gòu)造了基于打分的構(gòu)象生成模型。GeoDiff 則借助去噪擴(kuò)散概率模型 (Denoising Diffusion Probabilistic Model) 并且基于具有等變保證的 GNN 來(lái)構(gòu)造模型。等變流(Equivariant Flow)驗(yàn)證了基于等變核的標(biāo)準(zhǔn)化流(Normalizing Flow)的可行性。
對(duì)于點(diǎn)云的建模點(diǎn)云是對(duì)象的一種表示格式,它通過(guò)一組分配有坐標(biāo)的點(diǎn)來(lái)描述形狀。在對(duì)點(diǎn)云建模這一領(lǐng)域,包含一些經(jīng)典的數(shù)據(jù)和任務(wù)。ModelNet40 和 ScanObjectNN 是兩個(gè)經(jīng)典的點(diǎn)云數(shù)據(jù)集,其任務(wù)是做物品分類。因?yàn)樵邳c(diǎn)云數(shù)據(jù)里面不存在顯式的點(diǎn)和點(diǎn)的鏈接,在等變圖神經(jīng)網(wǎng)絡(luò)建模點(diǎn)云的時(shí)候,往往會(huì)以一個(gè)距離 d 作為閾值來(lái)構(gòu)造點(diǎn)和點(diǎn)之間的邊。TFN 和 SE(3)-Transformer 都在點(diǎn)云數(shù)據(jù)取得了相對(duì)于傳統(tǒng)方法具有競(jìng)爭(zhēng)力的性能。 未來(lái)展望 在對(duì)現(xiàn)有等變圖神經(jīng)網(wǎng)絡(luò)的方法和應(yīng)用進(jìn)行了系統(tǒng)性總結(jié)后。本綜述也對(duì)這一領(lǐng)域未來(lái)潛在發(fā)展方向進(jìn)行了一些討論:理論的完備性:不同于經(jīng)典圖神經(jīng)網(wǎng)絡(luò),等變神經(jīng)網(wǎng)絡(luò)還缺乏一系列的關(guān)于表達(dá)能力和泛化性的理論分析框架。現(xiàn)有的一些工作主要集中在討論消息傳播機(jī)制中存在的通用表達(dá)。但是對(duì)于模型的整體性質(zhì)依然不夠清晰。如何構(gòu)造一個(gè)完備的理論框架以指導(dǎo)模型的設(shè)計(jì)是未來(lái)十分有趣的方向。
大規(guī)模等變圖神經(jīng)網(wǎng)絡(luò),在前面已經(jīng)提到,基于群表示理論的方法都具有計(jì)算復(fù)雜度太高這一缺點(diǎn),這限制了等變圖神經(jīng)網(wǎng)絡(luò)在大規(guī)模數(shù)據(jù)上的應(yīng)用。尤其是在結(jié)合一些更復(fù)雜的結(jié)構(gòu),例注意力機(jī)制的時(shí)候,這一問(wèn)題會(huì)更加嚴(yán)重。如何有效的精簡(jiǎn)現(xiàn)有模型并且加速計(jì)算,使得現(xiàn)有的等變圖神經(jīng)網(wǎng)絡(luò)可以適用于大規(guī)模數(shù)據(jù),是一個(gè)重要方向。
對(duì)多層次結(jié)構(gòu)的建模:許多現(xiàn)實(shí)世界的系統(tǒng)都表現(xiàn)出復(fù)雜層次結(jié)構(gòu)。例如,有機(jī)分子由多個(gè)官能團(tuán)組成,蛋白質(zhì)由氨基酸組成。通過(guò)利用這些結(jié)構(gòu),我們可以設(shè)計(jì)出對(duì)多粒度多層次結(jié)構(gòu)的系統(tǒng)進(jìn)行建模。與現(xiàn)有的只有單層結(jié)構(gòu)的消息傳遞范式相比。這種層次結(jié)構(gòu)的等變模型可能可以更好的刻畫這種層次結(jié)構(gòu)信息,提高模型的性能和泛化性。
新的應(yīng)用和數(shù)據(jù):現(xiàn)有的等變圖神經(jīng)網(wǎng)絡(luò)模型大多數(shù)只是在規(guī)模和復(fù)雜性優(yōu)先的系統(tǒng)上進(jìn)行性能評(píng)估,例如模擬的 N-body 系統(tǒng)和小分子 MD 數(shù)據(jù)。未來(lái),我們需要在一些更具有挑戰(zhàn)性的任務(wù)上評(píng)估等變圖神經(jīng)網(wǎng)絡(luò)的有效性,這些挑戰(zhàn)包括:更多數(shù)量的對(duì)象,更復(fù)雜的交互,更多樣化的約束等等。近期以來(lái),在蛋白質(zhì)建模上的工作是一些有益的嘗試。但是由于數(shù)據(jù)收集的困難和數(shù)據(jù)質(zhì)量的限制,尚未有一個(gè)全面的可以評(píng)估各種方法的數(shù)據(jù)集出現(xiàn)。在未來(lái)如何將等變神經(jīng)網(wǎng)絡(luò)擴(kuò)展到更多,更復(fù)雜的領(lǐng)域,去解決現(xiàn)實(shí)問(wèn)題是一個(gè)很有意義的方向。