close
選自Github

作者:Chaitanya K. Joshi

機器之心編譯

機器之心編輯部

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

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

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

具體內容如下:

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

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

內存限制

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


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

硬件限制

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


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

可靠性限制

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

如何設計同時保留全圖語義和梯度信息的採樣程序?

處理巨型圖

二次採樣技術

現有論文在嘗試將巨型圖放入 GNN 時,關注點在於圖的子採樣,以將大圖拆分為可管理的子圖。ClusterGCN 利用聚類算法將圖拆分為聚類子圖,並通過將每個簇作為單次的小批量數據來訓練 GNN。在實踐中這樣做效果很好,特別是對於同質圖,同質圖中集群通常形成具有相似標籤的有意義的集群。

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

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

歷史節點嵌入

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

資料來源:《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。

資料來源:《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 的推理延遲和內存需求。


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

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


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

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 模型的性能,同時易於在生產系統中部署。

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

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

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

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

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


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

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

結論與展望

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

巨型圖 - 內存限制

稀疏計算——硬件限制

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


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

1. 數據準備——通過歷史節點嵌入查找,實現從對大規模圖採樣到 CPU-GPU 中進行混合訓練。
2. 高效架構——用於擴展到巨型網絡的圖增強 MLP,以及用於對批量圖數據進行實時推理的高效圖卷積設計。
3. 學習範式——將量化感知訓練(低精度模型權重和激活)與知識蒸餾(使用富有表現力的教師模型將 GNN 改進地更加高效)相結合,以最大限度地提高推理延遲和性能。


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

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

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

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

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

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


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

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

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

©THE END

轉載請聯繫本公眾號獲得授權

投稿或尋求報道:content@jiqizhixin.com

arrow
arrow
    全站熱搜
    創作者介紹
    創作者 鑽石舞台 的頭像
    鑽石舞台

    鑽石舞台

    鑽石舞台 發表在 痞客邦 留言(0) 人氣()