天天看點

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

選自Github

作者:Chaitanya K. Joshi

機器之心編譯

機器之心編輯部

圖神經網絡在應用到現實世界時會面臨很多挑戰,比如記憶體限制、硬體限制、可靠性限制等。在這篇文章中,劍橋大學在讀博士生 Chaitanya K. Joshi 從資料準備、高效架構和學習範式三個方向綜述了研究者們在克服這些問題時取得的進展。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

用于高效和可擴充的圖形表示學習的工具箱。

本文旨在概述關于高效圖神經網絡和可擴充圖表示學習的關鍵思想,并将介紹資料準備、GNN 架構和學習範式方面的關鍵進展,這些最新進展讓圖神經網絡能夠擴充到現實世界,并應用于實時場景。

具體内容如下:

圖神經網絡面臨的現實挑戰

圖神經網絡作為一種新興的深度學習體系架構,它可以操作諸如圖、集合、3D 點雲等不規則資料結構。最近幾年,GNN 的發展跨越社交網絡、推薦系統、生物醫學發現等不同領域。但在實際應用中,建構 GNN 面臨以下挑戰:

記憶體限制

現實世界的網絡可能非常龐大和複雜,例如 Facebook 有近 30 億活躍賬戶,這些賬戶以點贊、評論、分享等不同方式進行互動,進而在以賬戶為節點構成的圖中創造出無數個邊。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

現實世界中的的圖網絡,例如記錄所有 Facebook 使用者表以及他們互動方式的圖網絡,可能非常龐大且難以處理,以至于可能無法将這種巨型圖網絡安裝到 GPU 記憶體中以訓練神經網絡。

硬體限制

圖本質上是一種稀疏對象,GNN 按理說應該利用其稀疏性來進行高效和可擴充的計算。但是這說起來容易做起來難,因為現代 GPU 旨在處理矩陣上的密集運算。雖然針對稀疏矩陣的定制硬體加速器可以顯著提高 GNN 的及時性和可擴充性,但如何設計仍然是一個懸而未決的問題。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

現代 GPU 更适用于密集矩陣運算,而圖本質上是稀疏結構。除非鄰接矩陣非常稀疏,否則在實作 GNN 的過程中,将圖簡單地視為密集矩陣并使用掩碼來識别非連通節點通常更快。

可靠性限制

處理巨型圖的一種直接方法是将它們分成更小的子圖,并通過小批量梯度法下降訓練 GNN(每個子圖可以作為一個小批量資料)。但是,這種方法最大的問題在于:與樣本獨立的機器學習标準資料集不同,網絡資料的關系結構會在樣本之間産生統計依賴性。是以,要確定子圖保留完整圖的語義以及為訓練 GNN 提供可靠的梯度并不是一件簡單的事情。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

如何設計同時保留全圖語義和梯度資訊的采樣程式?

處理巨型圖

二次采樣技術

現有論文在嘗試将巨型圖放入 GNN 時,關注點在于圖的子采樣,以将大圖拆分為可管理的子圖。ClusterGCN 利用聚類算法将圖拆分為聚類子圖,并通過将每個簇作為單次的小批量資料來訓練 GNN。在實踐中這樣做效果很好,特别是對于同質圖,同質圖中叢集通常形成具有相似标簽的有意義的叢集。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《GraphSAINT: Graph Sampling Based Inductive Learning Method》

GraphSAINT 提出了一種更通用的機率圖采樣器來建構小批量子圖。可能的采樣方案包括統一節點 / 邊緣采樣以及随機遊走采樣。然而,由于上一節中強調的可靠性問題(語義和梯度資訊),與在全圖上訓練相比,子采樣方法可能會限制模型的性能。

曆史節點嵌入

GNNAutoScale (GAS) 是基本子采樣技術的一種很有前景的替代方案,用于将 GNN 應用到大型圖。GAS 建立在 Chen 等人之前關于曆史節點嵌入的工作之上,即将之前訓練得到的節點嵌入,重新用于現在的訓練。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings.》。

GAS 架構有兩個主要組成部分:首先,第一部分是建構一個小批量節點(執行快速随機子采樣)并修剪 GNN 計算圖以僅保留小批量内的節點及其 1 跳鄰居節點——這意味着 GAS 的尺度獨立于 GNN 深度。其次,每當 GNN 聚合需要小批量節點嵌入時,GAS 就會從存儲在 CPU 上的曆史嵌入中檢索它們。同時,目前小批量節點的曆史嵌入也不斷更新。

第二部分是與子采樣的關鍵差別——能夠使 GNN 最大限度地表達資訊,并将目前的小批量資料和曆史嵌入組合起來,得到完整的鄰域資訊并加以利用,同時確定對大型圖的可擴充性。

GAS 的作者還将他們的想法整合到流行的 PyTorch 幾何庫中。于是可以在非常大的圖上訓練大多數的消息傳遞 GNN,同時降低 GPU 記憶體需求并保持接近全批次的性能(即在全圖上訓練時的性能)。

工具包中用于擴充到大型圖的其他一些想法還包括:

[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.

可擴充且資源高效的 GNN 架構

圖增強 MLP

在開發可擴充 GNN 的時候,一個違反直覺的想法是:隻需在小批量節點上運作簡單的 MLP,而不考慮圖的關系結構!

Wu 等人的簡化圖卷積網絡(SGC)是第一個提出這個想法的工作。SGC 本質上是由 Kipf 和 Welling 通過将(昂貴但非學習的)稀疏鄰域特征聚合步驟與(廉價且可學習的)線性投影解耦,然後用 ReLU 非線性步驟來 “解構” 普通 GCN。在處理大型圖時,可以在 CPU 上高效地預先計算特征聚合(CPU 在處理稀疏操作方面表現不錯),然後可以對 “結構增強” 節點特征進行批處理并傳遞給在 GPU 上訓練的 MLP。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《Simplifying Graph Convolutional Networks》

SIGN:Rossi 等人的 Scalable Inception Graph Neural Networks 試圖通過從計算機視覺的 Inception 網絡中尋求靈感,并運作多個預計算步驟來為每個圖節點獲得更好的初始結構特征,進而将 SGC 的想法更進一步。

這一系列架構的其他有趣發展包括 Chen 等人探索結構增強 MLP 的理論局限性的工作,以及 Huang 等人的論文(該論文顯示,在節點特征上運作 MLP 後類似的标簽傳播方法的性能優于在同質圖上更繁瑣的 GNN)。

高效的圖卷積層

不幸的是,同質圖上的節點分類之外的許多任務可能需要比 SGC/SIGN 類模型更具表現力的 GNN 架構。這通常發生在對圖做分類或推理任務時,其中表現最好的 GNN 通常利用節點和邊的特征。

此類模型遵循 GNN 架構設計的消息傳遞風格(由 Petar Veli kovi 推廣),并且可以被認為是「各向異性的」,因為它們對每條邊進行了不同的處理(而普通 GCN 是「各向同性的」,因為相同的可學習權重被應用于每條邊)。然而,在每層的節點和邊上維護隐空間嵌入會顯著增加 GNN 的推理延遲和記憶體需求。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《Do We Need Anisotropic Graph Neural Networks?》

Tailor 等人的高效圖卷積 (EGC) 試圖解決這個困境。他們從基本的 GCN 設計開始(對于批量中等大小的圖具有良好的可擴充性),并設計了一個最大表達的卷積 GNN 版本,同時保留了 GCN 的可擴充性。令人印象深刻的是,它們在 Open Graph Benchmark 的圖分類任務中優于更複雜和繁瑣的基準模型。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

EGC 層也已內建到 PyTorch Geometric 中,可以作為即插即用的替代品來提高 GNN 的性能和可擴充性。

Li 等人的另一個有趣的想法是利用計算機視覺中高效的 ConvNet(可逆連接配接、組卷積、權重綁定和平衡模型)來提高 GNN 的記憶體和參數效率。他們的架構能夠訓練具有(前所未有的)1000 多個層的 GNN,并在 Open Graph Benchmark 的大規模節點分類任務中表現出色。

GNN 壓縮的學習範式

除了資料準備技術和有效的模型架構之外,學習模式,即模型的訓練方式,也可以顯著提高 GNN 的性能,并且降低延遲。

知識蒸餾可以提升性能

知識蒸餾(KD)是一種通用的神經網絡學習範式,它将知識從高性能但資源密集型的教師模型轉移到資源高效的學生身上。KD 的概念最初是由 Hinton 等人提出的,KD 訓練學生以比對教師模型的輸出 logits 以及标準的監督學習損失。

楊等人最近的工作以及随後發表的 Zhang 等人的工作,展示了這種簡單的基于 logit 的 KD 理念對于開發資源高效的 GNN 的強大功能:他們将富有表現力的 GNN 訓練為教師模型,并使用 KD 将知識轉移給 MLP 學生,以便在節點特征和圖結構高度相關的情況下更容易部署。這個想法可以擴充到上一節中的 SGC 或 SIGN 等圖形增強型 MLP,圖形增強型 MLP 可以顯着提高類 MLP 模型的性能,同時易于在生産系統中部署。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《Graph-less Neural Networks: Teaching Old MLPs New Tricks via Distillation》

在計算機視覺中,實踐者試圖超越基于 logit 的 KD,通過對齊潛在嵌入空間的損失函數将表示性知識從教師轉移到學生。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:《On Representation Knowledge Distillation for Graph Neural Networks》

Yang 等人的開創性工作首先通過訓練學生從教師的節點嵌入空間中保留局部拓撲結構來探索 GNN 的表示蒸餾。他們将這種方法稱為局部結構保留 (LSP),因為它鼓勵學生模仿教師節點嵌入空間中存在的直接鄰居的成對相似性。

最近,Joshi 等人擴充了 LSP 研究,通過保留直接鄰居之外的潛在互相作用來考慮全局拓撲結構。他們為 GNN 提出了兩個新的表示蒸餾目标:(1)一種顯式方法——全局結構保留,它擴充了 LSP 以考慮所有成對的相似性;(2) 一種隐式方法——圖形對比表示蒸餾,它使用對比學習将學生節點嵌入與教師節點嵌入在共享表示空間中對齊。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

Joshi 等人在 Open Graph Benchmark 資料集上的實驗表明,使用表示蒸餾訓練輕量級 GNN 可以顯著提高其實驗性能以及對噪聲或分布外資料的魯棒性。

GNN 的其他一些 KD 方法包括無教師蒸餾,也稱為自蒸餾,以及無資料蒸餾。

低精度的 GNN 的量化

量化感覺訓練(QAT)是另一種通用的神經網絡學習範式。雖然傳統的神經網絡模型權重和激活存儲為 32 位浮點數 FP32,但 QAT 訓練具有較低精度、整數權重和激活的模型,例如 INT8 或 INT4。低精度模型在推理延遲方面享有顯著優勢,盡管以降低性能為代價。

Tailor 等人的 DegreeQuant 提出了一種專門用于 GNN 的 QAT 技術。為計算機視覺 CNN 設計的通用 QAT 在應用于 GNN 時通常會導緻量化後的性能非常差。DegreeQuant 旨在通過巧妙地将基礎資料的圖結構整合到量化過程中來緩解這一問題:他們表明,具有許多鄰居(度數較高)的節點會導緻 QAT 期間的不穩定,并建議在執行 QAT 時随機屏蔽度數較高的節點。與 FP32 模型相比,這為 GNN 提供了更穩定的 QAT,并最大限度地減少了 INT8 的性能下降。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

量化 GNN 的其他想法包括利用 Zhao 等人的 Neural Architecture Search 和 Bahri 等人的二值化方法(這是量化的極端情況)。一般來說,QAT 的一些性能降低可以通過知識蒸餾來恢複,如 Bahri 等人以及 Joshi 等人在上一節的 Graph Contrastive Representation Distillation 論文中所示。

結論與展望

本文重點介紹高效的圖神經網絡和可擴充的圖表示學習。我們首先确定了現實世界 GNN 的理論和工程挑戰:

巨型圖 - 記憶體限制

稀疏計算——硬體限制

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

然後,我們介紹了三個關鍵思想,它們可能會成為開發高效且可擴充的 GNN 的工具:

1. 資料準備——通過曆史節點嵌入查找,實作從對大規模圖采樣到 CPU-GPU 中進行混合訓練。

2. 高效架構——用于擴充到巨型網絡的圖增強 MLP,以及用于對批量圖資料進行實時推理的高效圖卷積設計。

3. 學習範式——将量化感覺訓練(低精度模型權重和激活)與知識蒸餾(使用富有表現力的教師模型将 GNN 改進地更加高效)相結合,以最大限度地提高推理延遲和性能。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

在不久的将來,預計研究社群将繼續推進 GNN 網絡的高效化、可擴充性工具箱,并可能通過直接內建的方式出現在 PyTorch Geometric 和 DGL 等 GNN 庫中。我們還希望聽到越來越多的 GNN 處理真實世界圖和實時應用程式的成功案例。

從長遠來看,我們希望圖資料 + GNN 從一個深奧的新興研究領域轉變為用于機器學習研究和應用的标準資料 + 模型範式(很像 2D 圖像 + CNN,或文本 + Transformers)。是以,我們可能期望看到 GNN 更深入地內建到 PyTorch 或 TensorFlow 等标準架構中,為 GNN 開發專門的硬體加速器,以及更複雜的圖資料軟硬體協同設計。

事實上,這些努力可能已經在從圖資料和 GNN 中獲得巨大商業價值的公司中進行!

如需更深入地了解本文所涵蓋的主題,請參閱以下研究:

Abadal 等人的廣泛調查涵蓋了從 GNN 的基本原理到用于圖表示學習的硬體加速器設計的所有内容(本文未涉及)。

GNN落地不再難,一文總結高效GNN和可擴充圖表示學習最新進展

資料來源:圖神經網絡加速調查:算法視角。

文中提到的論文請參考原文。

繼續閱讀