GNN落地不再難,一文總結(jié)高效GNN和可擴(kuò)展圖表示學(xué)習(xí)最新進(jìn)展
圖神經(jīng)網(wǎng)絡(luò)在應(yīng)用到現(xiàn)實(shí)世界時(shí)會(huì)面臨很多挑戰(zhàn),比如內(nèi)存限制、硬件限制、可靠性限制等。在這篇文章中,劍橋大學(xué)在讀博士生 Chaitanya K. Joshi 從數(shù)據(jù)準(zhǔn)備、高效架構(gòu)和學(xué)習(xí)范式三個(gè)方向綜述了研究者們?cè)诳朔@些問(wèn)題時(shí)取得的進(jìn)展。
用于高效和可擴(kuò)展的圖形表示學(xué)習(xí)的工具箱。
本文旨在概述關(guān)于高效圖神經(jīng)網(wǎng)絡(luò)和可擴(kuò)展圖表示學(xué)習(xí)的關(guān)鍵思想,并將介紹數(shù)據(jù)準(zhǔn)備、GNN 架構(gòu)和學(xué)習(xí)范式方面的關(guān)鍵進(jìn)展,這些最新進(jìn)展讓圖神經(jīng)網(wǎng)絡(luò)能夠擴(kuò)展到現(xiàn)實(shí)世界,并應(yīng)用于實(shí)時(shí)場(chǎng)景。
具體內(nèi)容如下:
圖神經(jīng)網(wǎng)絡(luò)面臨的現(xiàn)實(shí)挑戰(zhàn)
圖神經(jīng)網(wǎng)絡(luò)作為一種新興的深度學(xué)習(xí)體系架構(gòu),它可以操作諸如圖、集合、3D 點(diǎn)云等不規(guī)則數(shù)據(jù)結(jié)構(gòu)。最近幾年,GNN 的發(fā)展跨越社交網(wǎng)絡(luò)、推薦系統(tǒng)、生物醫(yī)學(xué)發(fā)現(xiàn)等不同領(lǐng)域。但在實(shí)際應(yīng)用中,構(gòu)建 GNN 面臨以下挑戰(zhàn):
內(nèi)存限制
現(xiàn)實(shí)世界的網(wǎng)絡(luò)可能非常龐大和復(fù)雜,例如 Facebook 有近 30 億活躍賬戶,這些賬戶以點(diǎn)贊、評(píng)論、分享等不同方式進(jìn)行互動(dòng),從而在以賬戶為節(jié)點(diǎn)構(gòu)成的圖中創(chuàng)造出無(wú)數(shù)個(gè)邊。
現(xiàn)實(shí)世界中的的圖網(wǎng)絡(luò),例如記錄所有 Facebook 用戶表以及他們交互方式的圖網(wǎng)絡(luò),可能非常龐大且難以處理,以至于可能無(wú)法將這種巨型圖網(wǎng)絡(luò)安裝到 GPU 內(nèi)存中以訓(xùn)練神經(jīng)網(wǎng)絡(luò)。
硬件限制
圖本質(zhì)上是一種稀疏對(duì)象,GNN 按理說(shuō)應(yīng)該利用其稀疏性來(lái)進(jìn)行高效和可擴(kuò)展的計(jì)算。但是這說(shuō)起來(lái)容易做起來(lái)難,因?yàn)楝F(xiàn)代 GPU 旨在處理矩陣上的密集運(yùn)算。雖然針對(duì)稀疏矩陣的定制硬件加速器可以顯著提高 GNN 的及時(shí)性和可擴(kuò)展性,但如何設(shè)計(jì)仍然是一個(gè)懸而未決的問(wèn)題。
現(xiàn)代 GPU 更適用于密集矩陣運(yùn)算,而圖本質(zhì)上是稀疏結(jié)構(gòu)。除非鄰接矩陣非常稀疏,否則在實(shí)現(xiàn) GNN 的過(guò)程中,將圖簡(jiǎn)單地視為密集矩陣并使用掩碼來(lái)識(shí)別非連通節(jié)點(diǎn)通常更快。
可靠性限制
處理巨型圖的一種直接方法是將它們分成更小的子圖,并通過(guò)小批量梯度法下降訓(xùn)練 GNN(每個(gè)子圖可以作為一個(gè)小批量數(shù)據(jù))。但是,這種方法最大的問(wèn)題在于:與樣本獨(dú)立的機(jī)器學(xué)習(xí)標(biāo)準(zhǔn)數(shù)據(jù)集不同,網(wǎng)絡(luò)數(shù)據(jù)的關(guān)系結(jié)構(gòu)會(huì)在樣本之間產(chǎn)生統(tǒng)計(jì)依賴性。因此,要確保子圖保留完整圖的語(yǔ)義以及為訓(xùn)練 GNN 提供可靠的梯度并不是一件簡(jiǎn)單的事情。
如何設(shè)計(jì)同時(shí)保留全圖語(yǔ)義和梯度信息的采樣程序?
處理巨型圖
二次采樣技術(shù)
現(xiàn)有論文在嘗試將巨型圖放入 GNN 時(shí),關(guān)注點(diǎn)在于圖的子采樣,以將大圖拆分為可管理的子圖。ClusterGCN 利用聚類算法將圖拆分為聚類子圖,并通過(guò)將每個(gè)簇作為單次的小批量數(shù)據(jù)來(lái)訓(xùn)練 GNN。在實(shí)踐中這樣做效果很好,特別是對(duì)于同質(zhì)圖,同質(zhì)圖中集群通常形成具有相似標(biāo)簽的有意義的集群。
資料來(lái)源:《GraphSAINT: Graph Sampling Based Inductive Learning Method》
GraphSAINT 提出了一種更通用的概率圖采樣器來(lái)構(gòu)建小批量子圖??赡艿牟蓸臃桨赴ńy(tǒng)一節(jié)點(diǎn) / 邊緣采樣以及隨機(jī)游走采樣。然而,由于上一節(jié)中強(qiáng)調(diào)的可靠性問(wèn)題(語(yǔ)義和梯度信息),與在全圖上訓(xùn)練相比,子采樣方法可能會(huì)限制模型的性能。
歷史節(jié)點(diǎn)嵌入
GNNAutoScale (GAS) 是基本子采樣技術(shù)的一種很有前景的替代方案,用于將 GNN 應(yīng)用到大型圖。GAS 建立在 Chen 等人之前關(guān)于歷史節(jié)點(diǎn)嵌入的工作之上,即將之前訓(xùn)練得到的節(jié)點(diǎn)嵌入,重新用于現(xiàn)在的訓(xùn)練。
資料來(lái)源:《GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings.》。
GAS 框架有兩個(gè)主要組成部分:首先,第一部分是構(gòu)建一個(gè)小批量節(jié)點(diǎn)(執(zhí)行快速隨機(jī)子采樣)并修剪 GNN 計(jì)算圖以僅保留小批量?jī)?nèi)的節(jié)點(diǎn)及其 1 跳鄰居節(jié)點(diǎn)——這意味著 GAS 的尺度獨(dú)立于 GNN 深度。其次,每當(dāng) GNN 聚合需要小批量節(jié)點(diǎn)嵌入時(shí),GAS 就會(huì)從存儲(chǔ)在 CPU 上的歷史嵌入中檢索它們。同時(shí),當(dāng)前小批量節(jié)點(diǎn)的歷史嵌入也不斷更新。
第二部分是與子采樣的關(guān)鍵區(qū)別——能夠使 GNN 最大限度地表達(dá)信息,并將當(dāng)前的小批量數(shù)據(jù)和歷史嵌入組合起來(lái),得到完整的鄰域信息并加以利用,同時(shí)確保對(duì)大型圖的可擴(kuò)展性。
GAS 的作者還將他們的想法整合到流行的 PyTorch 幾何庫(kù)中。于是可以在非常大的圖上訓(xùn)練大多數(shù)的消息傳遞 GNN,同時(shí)降低 GPU 內(nèi)存需求并保持接近全批次的性能(即在全圖上訓(xùn)練時(shí)的性能)。
工具包中用于擴(kuò)展到大型圖的其他一些想法還包括:
[CVPR 2020] L2-GCN: Layer-Wise and Learned Efficient Training of Graph Convolutional Networks. Yuning You, Tianlong Chen, Zhangyang Wang, Yang Shen.
[KDD 2020] Scaling Graph Neural Networks with Approximate PageRank. Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberczki, Michal Lukasik, Stephan Günnemann.
[ICLR 2021] Graph Traversal with Tensor Functionals: A Meta-Algorithm for Scalable Learning. Elan Markowitz, Keshav Balasubramanian, Mehrnoosh Mirtaheri, Sami Abu-El-Haija, Bryan Perozzi, Greg Ver Steeg, Aram Galstyan.
[NeurIPS 2021] Decoupling the Depth and Scope of Graph Neural Networks. Hanqing Zeng, Muhan Zhang, Yinglong Xia, Ajitesh Srivastava, Andrey Malevich, Rajgopal Kannan, Viktor Prasanna, Long Jin, Ren Chen.
可擴(kuò)展且資源高效的 GNN 架構(gòu)
圖增強(qiáng) MLP
在開(kāi)發(fā)可擴(kuò)展 GNN 的時(shí)候,一個(gè)違反直覺(jué)的想法是:只需在小批量節(jié)點(diǎn)上運(yùn)行簡(jiǎn)單的 MLP,而不考慮圖的關(guān)系結(jié)構(gòu)!
Wu 等人的簡(jiǎn)化圖卷積網(wǎng)絡(luò)(SGC)是第一個(gè)提出這個(gè)想法的工作。SGC 本質(zhì)上是由 Kipf 和 Welling 通過(guò)將(昂貴但非學(xué)習(xí)的)稀疏鄰域特征聚合步驟與(廉價(jià)且可學(xué)習(xí)的)線性投影解耦,然后用 ReLU 非線性步驟來(lái) “解構(gòu)” 普通 GCN。在處理大型圖時(shí),可以在 CPU 上高效地預(yù)先計(jì)算特征聚合(CPU 在處理稀疏操作方面表現(xiàn)不錯(cuò)),然后可以對(duì) “結(jié)構(gòu)增強(qiáng)” 節(jié)點(diǎn)特征進(jìn)行批處理并傳遞給在 GPU 上訓(xùn)練的 MLP。
資料來(lái)源:《Simplifying Graph Convolutional Networks》
SIGN:Rossi 等人的 Scalable Inception Graph Neural Networks 試圖通過(guò)從計(jì)算機(jī)視覺(jué)的 Inception 網(wǎng)絡(luò)中尋求靈感,并運(yùn)行多個(gè)預(yù)計(jì)算步驟來(lái)為每個(gè)圖節(jié)點(diǎn)獲得更好的初始結(jié)構(gòu)特征,從而將 SGC 的想法更進(jìn)一步。
這一系列架構(gòu)的其他有趣發(fā)展包括 Chen 等人探索結(jié)構(gòu)增強(qiáng) MLP 的理論局限性的工作,以及 Huang 等人的論文(該論文顯示,在節(jié)點(diǎn)特征上運(yùn)行 MLP 后類似的標(biāo)簽傳播方法的性能優(yōu)于在同質(zhì)圖上更繁瑣的 GNN)。
高效的圖卷積層
不幸的是,同質(zhì)圖上的節(jié)點(diǎn)分類之外的許多任務(wù)可能需要比 SGC/SIGN 類模型更具表現(xiàn)力的 GNN 架構(gòu)。這通常發(fā)生在對(duì)圖做分類或推理任務(wù)時(shí),其中表現(xiàn)最好的 GNN 通常利用節(jié)點(diǎn)和邊的特征。
此類模型遵循 GNN 架構(gòu)設(shè)計(jì)的消息傳遞風(fēng)格(由 Petar Veli?kovi? 推廣),并且可以被認(rèn)為是「各向異性的」,因?yàn)樗鼈儗?duì)每條邊進(jìn)行了不同的處理(而普通 GCN 是「各向同性的」,因?yàn)橄嗤目蓪W(xué)習(xí)權(quán)重被應(yīng)用于每條邊)。然而,在每層的節(jié)點(diǎn)和邊上維護(hù)隱空間嵌入會(huì)顯著增加 GNN 的推理延遲和內(nèi)存需求。
資料來(lái)源:《Do We Need Anisotropic Graph Neural Networks?》
Tailor 等人的高效圖卷積 (EGC) 試圖解決這個(gè)困境。他們從基本的 GCN 設(shè)計(jì)開(kāi)始(對(duì)于批量中等大小的圖具有良好的可擴(kuò)展性),并設(shè)計(jì)了一個(gè)最大表達(dá)的卷積 GNN 版本,同時(shí)保留了 GCN 的可擴(kuò)展性。令人印象深刻的是,它們?cè)?Open Graph Benchmark 的圖分類任務(wù)中優(yōu)于更復(fù)雜和繁瑣的基準(zhǔn)模型。
資料來(lái)源:《Do We Need Anisotropic Graph Neural Networks?》
EGC 層也已集成到 PyTorch Geometric 中,可以作為即插即用的替代品來(lái)提高 GNN 的性能和可擴(kuò)展性。
Li 等人的另一個(gè)有趣的想法是利用計(jì)算機(jī)視覺(jué)中高效的 ConvNet(可逆連接、組卷積、權(quán)重綁定和平衡模型)來(lái)提高 GNN 的內(nèi)存和參數(shù)效率。他們的框架能夠訓(xùn)練具有(前所未有的)1000 多個(gè)層的 GNN,并在 Open Graph Benchmark 的大規(guī)模節(jié)點(diǎn)分類任務(wù)中表現(xiàn)出色。
GNN 壓縮的學(xué)習(xí)范式
除了數(shù)據(jù)準(zhǔn)備技術(shù)和有效的模型架構(gòu)之外,學(xué)習(xí)模式,即模型的訓(xùn)練方式,也可以顯著提高 GNN 的性能,并且降低延遲。
知識(shí)蒸餾可以提升性能
知識(shí)蒸餾(KD)是一種通用的神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)范式,它將知識(shí)從高性能但資源密集型的教師模型轉(zhuǎn)移到資源高效的學(xué)生身上。KD 的概念最初是由 Hinton 等人提出的,KD 訓(xùn)練學(xué)生以匹配教師模型的輸出 logits 以及標(biāo)準(zhǔn)的監(jiān)督學(xué)習(xí)損失。
楊等人最近的工作以及隨后發(fā)表的 Zhang 等人的工作,展示了這種簡(jiǎn)單的基于 logit 的 KD 理念對(duì)于開(kāi)發(fā)資源高效的 GNN 的強(qiáng)大功能:他們將富有表現(xiàn)力的 GNN 訓(xùn)練為教師模型,并使用 KD 將知識(shí)轉(zhuǎn)移給 MLP 學(xué)生,以便在節(jié)點(diǎn)特征和圖結(jié)構(gòu)高度相關(guān)的情況下更容易部署。這個(gè)想法可以擴(kuò)展到上一節(jié)中的 SGC 或 SIGN 等圖形增強(qiáng)型 MLP,圖形增強(qiáng)型 MLP 可以顯著提高類 MLP 模型的性能,同時(shí)易于在生產(chǎn)系統(tǒng)中部署。
資料來(lái)源:《Graph-less Neural Networks: Teaching Old MLPs New Tricks via Distillation》
在計(jì)算機(jī)視覺(jué)中,實(shí)踐者試圖超越基于 logit 的 KD,通過(guò)對(duì)齊潛在嵌入空間的損失函數(shù)將表示性知識(shí)從教師轉(zhuǎn)移到學(xué)生。
資料來(lái)源:《On Representation Knowledge Distillation for Graph Neural Networks》
Yang 等人的開(kāi)創(chuàng)性工作首先通過(guò)訓(xùn)練學(xué)生從教師的節(jié)點(diǎn)嵌入空間中保留局部拓?fù)浣Y(jié)構(gòu)來(lái)探索 GNN 的表示蒸餾。他們將這種方法稱為局部結(jié)構(gòu)保留 (LSP),因?yàn)樗膭?lì)學(xué)生模仿教師節(jié)點(diǎn)嵌入空間中存在的直接鄰居的成對(duì)相似性。
最近,Joshi 等人擴(kuò)展了 LSP 研究,通過(guò)保留直接鄰居之外的潛在相互作用來(lái)考慮全局拓?fù)浣Y(jié)構(gòu)。他們?yōu)?GNN 提出了兩個(gè)新的表示蒸餾目標(biāo):(1)一種顯式方法——全局結(jié)構(gòu)保留,它擴(kuò)展了 LSP 以考慮所有成對(duì)的相似性;(2) 一種隱式方法——圖形對(duì)比表示蒸餾,它使用對(duì)比學(xué)習(xí)將學(xué)生節(jié)點(diǎn)嵌入與教師節(jié)點(diǎn)嵌入在共享表示空間中對(duì)齊。
資料來(lái)源:《On Representation Knowledge Distillation for Graph Neural Networks》
Joshi 等人在 Open Graph Benchmark 數(shù)據(jù)集上的實(shí)驗(yàn)表明,使用表示蒸餾訓(xùn)練輕量級(jí) GNN 可以顯著提高其實(shí)驗(yàn)性能以及對(duì)噪聲或分布外數(shù)據(jù)的魯棒性。
GNN 的其他一些 KD 方法包括無(wú)教師蒸餾,也稱為自蒸餾,以及無(wú)數(shù)據(jù)蒸餾。
低精度的 GNN 的量化
量化感知訓(xùn)練(QAT)是另一種通用的神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)范式。雖然傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)模型權(quán)重和激活存儲(chǔ)為 32 位浮點(diǎn)數(shù) FP32,但 QAT 訓(xùn)練具有較低精度、整數(shù)權(quán)重和激活的模型,例如 INT8 或 INT4。低精度模型在推理延遲方面享有顯著優(yōu)勢(shì),盡管以降低性能為代價(jià)。
Tailor 等人的 DegreeQuant 提出了一種專門用于 GNN 的 QAT 技術(shù)。為計(jì)算機(jī)視覺(jué) CNN 設(shè)計(jì)的通用 QAT 在應(yīng)用于 GNN 時(shí)通常會(huì)導(dǎo)致量化后的性能非常差。DegreeQuant 旨在通過(guò)巧妙地將基礎(chǔ)數(shù)據(jù)的圖結(jié)構(gòu)整合到量化過(guò)程中來(lái)緩解這一問(wèn)題:他們表明,具有許多鄰居(度數(shù)較高)的節(jié)點(diǎn)會(huì)導(dǎo)致 QAT 期間的不穩(wěn)定,并建議在執(zhí)行 QAT 時(shí)隨機(jī)屏蔽度數(shù)較高的節(jié)點(diǎn)。與 FP32 模型相比,這為 GNN 提供了更穩(wěn)定的 QAT,并最大限度地減少了 INT8 的性能下降。
量化 GNN 的其他想法包括利用 Zhao 等人的 Neural Architecture Search 和 Bahri 等人的二值化方法(這是量化的極端情況)。一般來(lái)說(shuō),QAT 的一些性能降低可以通過(guò)知識(shí)蒸餾來(lái)恢復(fù),如 Bahri 等人以及 Joshi 等人在上一節(jié)的 Graph Contrastive Representation Distillation 論文中所示。
結(jié)論與展望
本文重點(diǎn)介紹高效的圖神經(jīng)網(wǎng)絡(luò)和可擴(kuò)展的圖表示學(xué)習(xí)。我們首先確定了現(xiàn)實(shí)世界 GNN 的理論和工程挑戰(zhàn):
巨型圖 - 內(nèi)存限制
稀疏計(jì)算——硬件限制
圖二次抽樣——可靠性限制
然后,我們介紹了三個(gè)關(guān)鍵思想,它們可能會(huì)成為開(kāi)發(fā)高效且可擴(kuò)展的 GNN 的工具:
1. 數(shù)據(jù)準(zhǔn)備——通過(guò)歷史節(jié)點(diǎn)嵌入查找,實(shí)現(xiàn)從對(duì)大規(guī)模圖采樣到 CPU-GPU 中進(jìn)行混合訓(xùn)練。2. 高效架構(gòu)——用于擴(kuò)展到巨型網(wǎng)絡(luò)的圖增強(qiáng) MLP,以及用于對(duì)批量圖數(shù)據(jù)進(jìn)行實(shí)時(shí)推理的高效圖卷積設(shè)計(jì)。3. 學(xué)習(xí)范式——將量化感知訓(xùn)練(低精度模型權(quán)重和激活)與知識(shí)蒸餾(使用富有表現(xiàn)力的教師模型將 GNN 改進(jìn)地更加高效)相結(jié)合,以最大限度地提高推理延遲和性能。
用于高效和可擴(kuò)展的圖形表示學(xué)習(xí)的工具箱。
在不久的將來(lái),預(yù)計(jì)研究社區(qū)將繼續(xù)推進(jìn) GNN 網(wǎng)絡(luò)的高效化、可擴(kuò)展性工具箱,并可能通過(guò)直接集成的方式出現(xiàn)在 PyTorch Geometric 和 DGL 等 GNN 庫(kù)中。我們還希望聽(tīng)到越來(lái)越多的 GNN 處理真實(shí)世界圖和實(shí)時(shí)應(yīng)用程序的成功案例。
從長(zhǎng)遠(yuǎn)來(lái)看,我們希望圖數(shù)據(jù) + GNN 從一個(gè)深?yuàn)W的新興研究領(lǐng)域轉(zhuǎn)變?yōu)橛糜跈C(jī)器學(xué)習(xí)研究和應(yīng)用的標(biāo)準(zhǔn)數(shù)據(jù) + 模型范式(很像 2D 圖像 + CNN,或文本 + Transformers)。因此,我們可能期望看到 GNN 更深入地集成到 PyTorch 或 TensorFlow 等標(biāo)準(zhǔn)框架中,為 GNN 開(kāi)發(fā)專門的硬件加速器,以及更復(fù)雜的圖數(shù)據(jù)軟硬件協(xié)同設(shè)計(jì)。
事實(shí)上,這些努力可能已經(jīng)在從圖數(shù)據(jù)和 GNN 中獲得巨大商業(yè)價(jià)值的公司中進(jìn)行!
如需更深入地了解本文所涵蓋的主題,請(qǐng)參閱以下研究:
Abadal 等人的廣泛調(diào)查涵蓋了從 GNN 的基本原理到用于圖表示學(xué)習(xí)的硬件加速器設(shè)計(jì)的所有內(nèi)容(本文未涉及)。
資料來(lái)源:圖神經(jīng)網(wǎng)絡(luò)加速調(diào)查:算法視角。
文中提到的論文請(qǐng)參考原文。
原文鏈接:https://www.chaitjo.com/post/efficient-gnns/
*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布,僅代表博主個(gè)人觀點(diǎn),如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。