GNN落地不再難，一文總結(jié)高效GNN和可擴展圖表示學(xué)習(xí)最新進展

圖神經(jīng)網(wǎng)絡(luò)在應(yīng)用到現(xiàn)實世界時會面臨很多挑戰(zhàn)，比如內(nèi)存限制、硬件限制、可靠性限制等。在這篇文章中，劍橋大學(xué)在讀博士生 Chaitanya K. Joshi 從數(shù)據(jù)準(zhǔn)備、高效架構(gòu)和學(xué)習(xí)范式三個方向綜述了研究者們在克服這些問題時取得的進展。

用于高效和可擴展的圖形表示學(xué)習(xí)的工具箱。
本文旨在概述關(guān)于高效圖神經(jīng)網(wǎng)絡(luò)和可擴展圖表示學(xué)習(xí)的關(guān)鍵思想，并將介紹數(shù)據(jù)準(zhǔn)備、GNN 架構(gòu)和學(xué)習(xí)范式方面的關(guān)鍵進展，這些最新進展讓圖神經(jīng)網(wǎng)絡(luò)能夠擴展到現(xiàn)實世界，并應(yīng)用于實時場景。
具體內(nèi)容如下：
圖神經(jīng)網(wǎng)絡(luò)面臨的現(xiàn)實挑戰(zhàn)
圖神經(jīng)網(wǎng)絡(luò)作為一種新興的深度學(xué)習(xí)體系架構(gòu)，它可以操作諸如圖、集合、3D 點云等不規(guī)則數(shù)據(jù)結(jié)構(gòu)。最近幾年，GNN 的發(fā)展跨越社交網(wǎng)絡(luò)、推薦系統(tǒng)、生物醫(yī)學(xué)發(fā)現(xiàn)等不同領(lǐng)域。但在實際應(yīng)用中，構(gòu)建 GNN 面臨以下挑戰(zhàn)：
內(nèi)存限制
現(xiàn)實世界的網(wǎng)絡(luò)可能非常龐大和復(fù)雜，例如 Facebook 有近 30 億活躍賬戶，這些賬戶以點贊、評論、分享等不同方式進行互動，從而在以賬戶為節(jié)點構(gòu)成的圖中創(chuàng)造出無數(shù)個邊。
現(xiàn)實世界中的的圖網(wǎng)絡(luò)，例如記錄所有 Facebook 用戶表以及他們交互方式的圖網(wǎng)絡(luò)，可能非常龐大且難以處理，以至于可能無法將這種巨型圖網(wǎng)絡(luò)安裝到 GPU 內(nèi)存中以訓(xùn)練神經(jīng)網(wǎng)絡(luò)。
硬件限制
圖本質(zhì)上是一種稀疏對象，GNN 按理說應(yīng)該利用其稀疏性來進行高效和可擴展的計算。但是這說起來容易做起來難，因為現(xiàn)代 GPU 旨在處理矩陣上的密集運算。雖然針對稀疏矩陣的定制硬件加速器可以顯著提高 GNN 的及時性和可擴展性，但如何設(shè)計仍然是一個懸而未決的問題。
現(xiàn)代 GPU 更適用于密集矩陣運算，而圖本質(zhì)上是稀疏結(jié)構(gòu)。除非鄰接矩陣非常稀疏，否則在實現(xiàn) GNN 的過程中，將圖簡單地視為密集矩陣并使用掩碼來識別非連通節(jié)點通常更快。
可靠性限制
處理巨型圖的一種直接方法是將它們分成更小的子圖，并通過小批量梯度法下降訓(xùn)練 GNN（每個子圖可以作為一個小批量數(shù)據(jù)）。但是，這種方法最大的問題在于：與樣本獨立的機器學(xué)習(xí)標(biāo)準(zhǔn)數(shù)據(jù)集不同，網(wǎng)絡(luò)數(shù)據(jù)的關(guān)系結(jié)構(gòu)會在樣本之間產(chǎn)生統(tǒng)計依賴性。因此，要確保子圖保留完整圖的語義以及為訓(xùn)練 GNN 提供可靠的梯度并不是一件簡單的事情。
如何設(shè)計同時保留全圖語義和梯度信息的采樣程序？
處理巨型圖
二次采樣技術(shù)
現(xiàn)有論文在嘗試將巨型圖放入 GNN 時，關(guān)注點在于圖的子采樣，以將大圖拆分為可管理的子圖。ClusterGCN 利用聚類算法將圖拆分為聚類子圖，并通過將每個簇作為單次的小批量數(shù)據(jù)來訓(xùn)練 GNN。在實踐中這樣做效果很好，特別是對于同質(zhì)圖，同質(zhì)圖中集群通常形成具有相似標(biāo)簽的有意義的集群。
資料來源：《GraphSAINT: Graph Sampling Based Inductive Learning Method》
GraphSAINT 提出了一種更通用的概率圖采樣器來構(gòu)建小批量子圖?？赡艿牟蓸臃桨赴ńy(tǒng)一節(jié)點 / 邊緣采樣以及隨機游走采樣。然而，由于上一節(jié)中強調(diào)的可靠性問題（語義和梯度信息），與在全圖上訓(xùn)練相比，子采樣方法可能會限制模型的性能。
歷史節(jié)點嵌入
GNNAutoScale (GAS) 是基本子采樣技術(shù)的一種很有前景的替代方案，用于將 GNN 應(yīng)用到大型圖。GAS 建立在 Chen 等人之前關(guān)于歷史節(jié)點嵌入的工作之上，即將之前訓(xùn)練得到的節(jié)點嵌入，重新用于現(xiàn)在的訓(xùn)練。
資料來源：《GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings.》。
GAS 框架有兩個主要組成部分：首先，第一部分是構(gòu)建一個小批量節(jié)點（執(zhí)行快速隨機子采樣）并修剪 GNN 計算圖以僅保留小批量內(nèi)的節(jié)點及其 1 跳鄰居節(jié)點——這意味著 GAS 的尺度獨立于 GNN 深度。其次，每當(dāng) GNN 聚合需要小批量節(jié)點嵌入時，GAS 就會從存儲在 CPU 上的歷史嵌入中檢索它們。同時，當(dāng)前小批量節(jié)點的歷史嵌入也不斷更新。
第二部分是與子采樣的關(guān)鍵區(qū)別——能夠使 GNN 最大限度地表達(dá)信息，并將當(dāng)前的小批量數(shù)據(jù)和歷史嵌入組合起來，得到完整的鄰域信息并加以利用，同時確保對大型圖的可擴展性。
GAS 的作者還將他們的想法整合到流行的 PyTorch 幾何庫中。于是可以在非常大的圖上訓(xùn)練大多數(shù)的消息傳遞 GNN，同時降低 GPU 內(nèi)存需求并保持接近全批次的性能（即在全圖上訓(xùn)練時的性能）。
工具包中用于擴展到大型圖的其他一些想法還包括：
[CVPR 2020] L²-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.

可擴展且資源高效的 GNN 架構(gòu)
圖增強 MLP
在開發(fā)可擴展 GNN 的時候，一個違反直覺的想法是：只需在小批量節(jié)點上運行簡單的 MLP，而不考慮圖的關(guān)系結(jié)構(gòu)！
Wu 等人的簡化圖卷積網(wǎng)絡(luò)（SGC）是第一個提出這個想法的工作。SGC 本質(zhì)上是由 Kipf 和 Welling 通過將（昂貴但非學(xué)習(xí)的）稀疏鄰域特征聚合步驟與（廉價且可學(xué)習(xí)的）線性投影解耦，然后用 ReLU 非線性步驟來 “解構(gòu)” 普通 GCN。在處理大型圖時，可以在 CPU 上高效地預(yù)先計算特征聚合（CPU 在處理稀疏操作方面表現(xiàn)不錯），然后可以對 “結(jié)構(gòu)增強” 節(jié)點特征進行批處理并傳遞給在 GPU 上訓(xùn)練的 MLP。
資料來源：《Simplifying Graph Convolutional Networks》
SIGN：Rossi 等人的 Scalable Inception Graph Neural Networks 試圖通過從計算機視覺的 Inception 網(wǎng)絡(luò)中尋求靈感，并運行多個預(yù)計算步驟來為每個圖節(jié)點獲得更好的初始結(jié)構(gòu)特征，從而將 SGC 的想法更進一步。
這一系列架構(gòu)的其他有趣發(fā)展包括 Chen 等人探索結(jié)構(gòu)增強 MLP 的理論局限性的工作，以及 Huang 等人的論文（該論文顯示，在節(jié)點特征上運行 MLP 后類似的標(biāo)簽傳播方法的性能優(yōu)于在同質(zhì)圖上更繁瑣的 GNN）。
高效的圖卷積層
不幸的是，同質(zhì)圖上的節(jié)點分類之外的許多任務(wù)可能需要比 SGC/SIGN 類模型更具表現(xiàn)力的 GNN 架構(gòu)。這通常發(fā)生在對圖做分類或推理任務(wù)時，其中表現(xiàn)最好的 GNN 通常利用節(jié)點和邊的特征。
此類模型遵循 GNN 架構(gòu)設(shè)計的消息傳遞風(fēng)格（由 Petar Veli?kovi? 推廣），并且可以被認(rèn)為是「各向異性的」，因為它們對每條邊進行了不同的處理（而普通 GCN 是「各向同性的」，因為相同的可學(xué)習(xí)權(quán)重被應(yīng)用于每條邊）。然而，在每層的節(jié)點和邊上維護隱空間嵌入會顯著增加 GNN 的推理延遲和內(nèi)存需求。
資料來源：《Do We Need Anisotropic Graph Neural Networks?》
Tailor 等人的高效圖卷積 (EGC) 試圖解決這個困境。他們從基本的 GCN 設(shè)計開始（對于批量中等大小的圖具有良好的可擴展性），并設(shè)計了一個最大表達(dá)的卷積 GNN 版本，同時保留了 GCN 的可擴展性。令人印象深刻的是，它們在 Open Graph Benchmark 的圖分類任務(wù)中優(yōu)于更復(fù)雜和繁瑣的基準(zhǔn)模型。
資料來源：《Do We Need Anisotropic Graph Neural Networks?》
EGC 層也已集成到 PyTorch Geometric 中，可以作為即插即用的替代品來提高 GNN 的性能和可擴展性。
Li 等人的另一個有趣的想法是利用計算機視覺中高效的 ConvNet（可逆連接、組卷積、權(quán)重綁定和平衡模型）來提高 GNN 的內(nèi)存和參數(shù)效率。他們的框架能夠訓(xùn)練具有（前所未有的）1000 多個層的 GNN，并在 Open Graph Benchmark 的大規(guī)模節(jié)點分類任務(wù)中表現(xiàn)出色。
GNN 壓縮的學(xué)習(xí)范式
除了數(shù)據(jù)準(zhǔn)備技術(shù)和有效的模型架構(gòu)之外，學(xué)習(xí)模式，即模型的訓(xùn)練方式，也可以顯著提高 GNN 的性能，并且降低延遲。
知識蒸餾可以提升性能
知識蒸餾（KD）是一種通用的神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)范式，它將知識從高性能但資源密集型的教師模型轉(zhuǎn)移到資源高效的學(xué)生身上。KD 的概念最初是由 Hinton 等人提出的，KD 訓(xùn)練學(xué)生以匹配教師模型的輸出 logits 以及標(biāo)準(zhǔn)的監(jiān)督學(xué)習(xí)損失。
楊等人最近的工作以及隨后發(fā)表的 Zhang 等人的工作，展示了這種簡單的基于 logit 的 KD 理念對于開發(fā)資源高效的 GNN 的強大功能：他們將富有表現(xiàn)力的 GNN 訓(xùn)練為教師模型，并使用 KD 將知識轉(zhuǎn)移給 MLP 學(xué)生，以便在節(jié)點特征和圖結(jié)構(gòu)高度相關(guān)的情況下更容易部署。這個想法可以擴展到上一節(jié)中的 SGC 或 SIGN 等圖形增強型 MLP，圖形增強型 MLP 可以顯著提高類 MLP 模型的性能，同時易于在生產(chǎn)系統(tǒng)中部署。
資料來源：《Graph-less Neural Networks: Teaching Old MLPs New Tricks via Distillation》
在計算機視覺中，實踐者試圖超越基于 logit 的 KD，通過對齊潛在嵌入空間的損失函數(shù)將表示性知識從教師轉(zhuǎn)移到學(xué)生。
資料來源：《On Representation Knowledge Distillation for Graph Neural Networks》
Yang 等人的開創(chuàng)性工作首先通過訓(xùn)練學(xué)生從教師的節(jié)點嵌入空間中保留局部拓?fù)浣Y(jié)構(gòu)來探索 GNN 的表示蒸餾。他們將這種方法稱為局部結(jié)構(gòu)保留 (LSP)，因為它鼓勵學(xué)生模仿教師節(jié)點嵌入空間中存在的直接鄰居的成對相似性。
最近，Joshi 等人擴展了 LSP 研究，通過保留直接鄰居之外的潛在相互作用來考慮全局拓?fù)浣Y(jié)構(gòu)。他們?yōu)?GNN 提出了兩個新的表示蒸餾目標(biāo)：（1）一種顯式方法——全局結(jié)構(gòu)保留，它擴展了 LSP 以考慮所有成對的相似性；(2) 一種隱式方法——圖形對比表示蒸餾，它使用對比學(xué)習(xí)將學(xué)生節(jié)點嵌入與教師節(jié)點嵌入在共享表示空間中對齊。
資料來源：《On Representation Knowledge Distillation for Graph Neural Networks》
Joshi 等人在 Open Graph Benchmark 數(shù)據(jù)集上的實驗表明，使用表示蒸餾訓(xùn)練輕量級 GNN 可以顯著提高其實驗性能以及對噪聲或分布外數(shù)據(jù)的魯棒性。
GNN 的其他一些 KD 方法包括無教師蒸餾，也稱為自蒸餾，以及無數(shù)據(jù)蒸餾。
低精度的 GNN 的量化
量化感知訓(xùn)練（QAT）是另一種通用的神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)范式。雖然傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)模型權(quán)重和激活存儲為 32 位浮點數(shù) FP32，但 QAT 訓(xùn)練具有較低精度、整數(shù)權(quán)重和激活的模型，例如 INT8 或 INT4。低精度模型在推理延遲方面享有顯著優(yōu)勢，盡管以降低性能為代價。
Tailor 等人的 DegreeQuant 提出了一種專門用于 GNN 的 QAT 技術(shù)。為計算機視覺 CNN 設(shè)計的通用 QAT 在應(yīng)用于 GNN 時通常會導(dǎo)致量化后的性能非常差。DegreeQuant 旨在通過巧妙地將基礎(chǔ)數(shù)據(jù)的圖結(jié)構(gòu)整合到量化過程中來緩解這一問題：他們表明，具有許多鄰居（度數(shù)較高）的節(jié)點會導(dǎo)致 QAT 期間的不穩(wěn)定，并建議在執(zhí)行 QAT 時隨機屏蔽度數(shù)較高的節(jié)點。與 FP32 模型相比，這為 GNN 提供了更穩(wěn)定的 QAT，并最大限度地減少了 INT8 的性能下降。

量化 GNN 的其他想法包括利用 Zhao 等人的 Neural Architecture Search 和 Bahri 等人的二值化方法（這是量化的極端情況）。一般來說，QAT 的一些性能降低可以通過知識蒸餾來恢復(fù)，如 Bahri 等人以及 Joshi 等人在上一節(jié)的 Graph Contrastive Representation Distillation 論文中所示。
結(jié)論與展望
本文重點介紹高效的圖神經(jīng)網(wǎng)絡(luò)和可擴展的圖表示學(xué)習(xí)。我們首先確定了現(xiàn)實世界 GNN 的理論和工程挑戰(zhàn)：

• 巨型圖 - 內(nèi)存限制

• 稀疏計算——硬件限制

• 圖二次抽樣——可靠性限制

然后，我們介紹了三個關(guān)鍵思想，它們可能會成為開發(fā)高效且可擴展的 GNN 的工具：
1. 數(shù)據(jù)準(zhǔn)備——通過歷史節(jié)點嵌入查找，實現(xiàn)從對大規(guī)模圖采樣到 CPU-GPU 中進行混合訓(xùn)練。2. 高效架構(gòu)——用于擴展到巨型網(wǎng)絡(luò)的圖增強 MLP，以及用于對批量圖數(shù)據(jù)進行實時推理的高效圖卷積設(shè)計。3. 學(xué)習(xí)范式——將量化感知訓(xùn)練（低精度模型權(quán)重和激活）與知識蒸餾（使用富有表現(xiàn)力的教師模型將 GNN 改進地更加高效）相結(jié)合，以最大限度地提高推理延遲和性能。
用于高效和可擴展的圖形表示學(xué)習(xí)的工具箱。
在不久的將來，預(yù)計研究社區(qū)將繼續(xù)推進 GNN 網(wǎng)絡(luò)的高效化、可擴展性工具箱，并可能通過直接集成的方式出現(xiàn)在 PyTorch Geometric 和 DGL 等 GNN 庫中。我們還希望聽到越來越多的 GNN 處理真實世界圖和實時應(yīng)用程序的成功案例。
從長遠(yuǎn)來看，我們希望圖數(shù)據(jù) + GNN 從一個深奧的新興研究領(lǐng)域轉(zhuǎn)變?yōu)橛糜跈C器學(xué)習(xí)研究和應(yīng)用的標(biāo)準(zhǔn)數(shù)據(jù) + 模型范式（很像 2D 圖像 + CNN，或文本 + Transformers）。因此，我們可能期望看到 GNN 更深入地集成到 PyTorch 或 TensorFlow 等標(biāo)準(zhǔn)框架中，為 GNN 開發(fā)專門的硬件加速器，以及更復(fù)雜的圖數(shù)據(jù)軟硬件協(xié)同設(shè)計。
事實上，這些努力可能已經(jīng)在從圖數(shù)據(jù)和 GNN 中獲得巨大商業(yè)價值的公司中進行！
如需更深入地了解本文所涵蓋的主題，請參閱以下研究：
Abadal 等人的廣泛調(diào)查涵蓋了從 GNN 的基本原理到用于圖表示學(xué)習(xí)的硬件加速器設(shè)計的所有內(nèi)容（本文未涉及）。
資料來源：圖神經(jīng)網(wǎng)絡(luò)加速調(diào)查：算法視角。
文中提到的論文請參考原文。

原文鏈接：https://www.chaitjo.com/post/efficient-gnns/