【GNN】Cluster-GCN：一個簡單又有效的 Trick

今天學習的是 Google 2019 年的工作《Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks》，發表於 KDD 2019。

目前，GCN 已經成功應用於各個領域，但是大規模的 GCN 訓練仍然是一個非常大的挑戰。目前基於 SGD 優化的 GCN 算法要麼面臨着隨着層數呈指數增長的計算成本，要麼面臨着保存整個圖和節點 Embedding 的巨大內存需求。

針對這個問題，作者提出了一種適用於基於 SGD 訓練新的圖聚類結構——Cluster-GCN。其核心思想是先利用圖聚類算法來區分子圖進行採樣，並限制該子圖中的鄰居搜索。這種策略雖然簡單，但是對提高內存和計算效率而言非常有效，同時也能夠保證算法的精度。

爲了測試算法的可擴展性，作者創建了一個新的 Amazon 數據集，比之前的 Reddit 大五倍，並在該數據集上取得了更快、更少內存。此外，隨着層數增加，Cluster-GCN 的預測精度也取得了 SOTA 的成績。

1.Introduction

由於節點的依賴性，GCN 的訓練需要消耗大量內存來計算圖上節點的 Embedding。

一些經典模型如 GCN 採用了 full-batch 的 SGD 優化算法，要計算整個梯度則需要存儲所有中間的 Embedding，因此，其是不可擴展的。此外，雖然每個 epoch 也只能更新一次參數。

GraphSAGE 中提出 mini-batch 的 SGD 優化方法，由於每次更新只基於一個 mini-batch，所以內存的需求降低，並在每個 epoch 中可以進行多次更新，從而收斂速度更快。然而，隨着層數加深，每個節點的感受野越來越大，其計算單個節點的計算開銷也會越來越大。針對這個問題，GraphSAGE 通過使用固定大小的鄰居採樣，同時 FastGCN 的重要性採樣可以一定程度上解決計算開銷，但是隨着 GCN 的深度加深，計算開銷問題依然沒法解決。

VR-GCN 提出利用方差來控制鄰居的採樣節點，儘管減少了採樣的大小，但是它需要將所有節點的中間 Embedding 存儲於內存中，導致其可擴展性較差。

下表展示了不同模型的時間複雜度和空間複雜度：

作者在實驗中發現 mini-batch 的算法效率與 batch 內節點與 batch 外節點間的連接數量成正比，針對這一現象，作者構造了節點的分區，使同一分區中的節點之間的圖連接於不同分區中的節點之間的圖連接更多。

在此基礎上，作者提出了基於圖聚類結構的新算法——Cluster GCN，並設計了一個隨機多聚類框架（ 「stochastic multi-clustering framework」 ）來提高 Cluster-GCN 的收斂性，從而降低了內存消耗和計算消耗。

2.Cluster-GCN

我們知道，基於 mini-batch 的 SGD 可以在單個 epoch 中更新多次，從而使得其比 full batch 具有更快的收斂速度，但是前者每個 epoch 所花的時間都更長。

出現這種情況主要是 SGD 在訓練時引入額外的計算開銷，我們簡單介紹下。

首先給出 SGD 的計算公式：

在計算節點 i 相關梯度時，需要節點 i 的 Embedding，而其計算需要依賴前一層鄰居的 Embedding，而前一層的節點的 Embedding 需要前前一層的鄰居節點的 Embedding，並如此嵌套下去，直至第一層。

考慮每個節點平均度數爲 d，GCN 網絡有 L 層，爲了獲取節點 i 相關的梯度需要對途中節點聚合個節點的特徵。也就是說需要獲取途中節點的 hop-k（k=1,..,L）鄰居的信息來執行一次更新。此外，還需要和權重矩陣相乘，所以每次計算 Embedding 還需要的時間，所以綜合起來平均計算一個節點相關的梯度的時間複雜度爲。

如果一個 batch 包含多個節點，那麼時間複雜度就沒那麼直觀了，因爲不同節點可能會有重疊的 hop-k 鄰居，並且 Embedding 計算的數量可以小於最壞情況。爲了反映 mini-batch SGD 的計算效率，作者定義了 Embedding utilization 的概念來表達計算效率。

如果節點 i 在第 l 層的 Embedding在計算第 l+1 層計算時被重用了 u 次，那麼就說的 Embedding utilization 爲 u。

對基於隨機抽樣的 mini-batch SGD 而言，由於圖通常比較大且稀疏，所以 u 通常非常小，所以 mini-batch SGD 需要計算每個 batch 的個 Embedding，這將導致每次更新時間複雜度爲，epoch 的時間複雜度爲。

full batch 梯度下降具有最大的 Embedding utilization——每個 Embedding 將在上層重複使用 d 次。因此，full-batch 的每個 epoch 只需要計算個 Embedding，這就意味着只需要計算個 Embedding 就可以計算一個節點的梯度。

針對這樣一種現象，作者爲了最大化 Embedding utilization，設計了 Cluster-GCN，旨在設計一個 batch 來最大化 batch 內邊的數量。

2.1 Vanilla Cluster-GCN

對於一個圖 G 而言，將其分爲 c 組：

其中，只包含中節點之間的邊。

對節點進行重組後，鄰接矩陣被劃分爲個子矩陣：

其中：

其中，爲子圖的內在鄰接矩陣；爲圖的鄰接矩陣；爲 A 的所有非對角塊組成的矩陣；爲和之間組成的鄰接矩陣。

類似的，也可以對特徵矩陣和訓練標籤按照子圖進行劃分。

損失函數也被分解爲：

Cluster-GCN 便是基於上面的公式，在每一步中，先對矩陣進行採樣，然後根據的梯度進行 SGD 更新，這裏只需要當前 batch 上的子圖的鄰接矩陣、特徵矩陣、標籤向量和權重矩陣。這相比於之前的 SGD 訓練所使用的鄰接採樣更容易實現，速度也更快。

作者使用 Metis 和 Craclus 等聚類算法在圖中的的頂點上構建分區，使簇內連接大於簇間鏈接，從而得到更好的聚類和社區結構。

劃分簇的意義在於：

對於每個 batch 而言，Embedding utilization 相當於簇內的連接。每個節點及其相鄰節點通常位於同一簇內，因此經過幾次後跳躍後，鄰接節點大概率還是在簇內；
利用來代替，誤差與簇間的的連接成正比，所以需要使得簇間的連接數量儘可能少。

下圖爲全圖和聚類分區圖：

下表爲兩種不同數據集的分區策略（隨機和 metis）及對應的訓練精度，可以看到聚類劃分還是很有必要的。

2.2 Stochastic Multiple Partitions

儘管 vanilla Cluster-GCN 能夠減少計算開銷和內存開銷，但仍然存在兩個問題：

圖被分割後，中的一些連接會被刪除，影響性能；
聚類後的分佈會與原始數據集有區別，從而導致 SGD 更新時有偏差。

下圖爲 Reddit 數據集中標籤分佈不平衡的案例，通過每個簇的標籤分佈計算其熵值，與隨機分割相比，可以清楚的看到聚類分區的簇的熵較小，這表明簇的標籤分佈偏向於某些特徵的標籤，所以這會增加不同 batch 的梯度更新的差異，並影響 SGD 的收斂性。

爲了解決這個問題，作者提出隨機多聚類方法（stochastic multiple clustering）對簇進行合併，從而減少 batch 間的差異。

作者首先將圖分爲多個小簇，然後隨機選擇 q 個簇併到 batch 中，這樣便可以減少 batch 之間的差異。

下圖展示了每個 epoch 隨機組合的 batch：

下圖兩種方式的對比，隨機多聚類方法的收斂速度更快：

2.3 Issues of training deeper GCN

作者提出了一個簡單的技術來改進深度 GCN 的訓練，核心思想在於放大每個 GCN 層中使用的鄰接矩陣 A 的對角部分，並通過這種方式在每個 GCN 層的聚合中對上一層的 Embedding 添加更多的權重：

但這種方法有些問題，比如這種方法無視相鄰節點的數量，而對所有節點使用相同的權重。此外，當層數增加時，其數值可能會呈現指數型爆炸。

所以作者先對鄰接矩陣進行標準化：

然後考慮：

這種新的標準化策略達到了 SOTA 的效果。

3.Experiment

簡單看一下實驗。

首先是數據集：

實驗所用參數：

不同數量的隱藏層下的模型內存消耗：

訓練時間和準確度：

在大數據集下實驗：

諸多模型的測試精度：

4.Conclusion

本文提出了一種新的訓練算法 Cluster-GCN，核心思想在於利用聚類算法將大圖劃分爲多個簇，劃分遵循簇間連接少而簇內連接多的原則，這種簡單的方法有效的減少了內存和計算資源的消耗，同時也能取得非常好的預測精度。

5.Reference

《Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks》
benedekrozemberczki/ClusterGCN
google-research/cluster_gcn

關於AINLP

AINLP 是一個有趣有AI的自然語言處理社區，專注於 AI、NLP、機器學習、深度學習、推薦算法等相關技術的分享，主題包括文本摘要、智能問答、聊天機器人、機器翻譯、自動生成、知識圖譜、預訓練模型、推薦系統、計算廣告、招聘信息、求職經驗分享等，歡迎關注！加技術交流羣請添加AINLPer(id：ainlper)，備註工作/研究方向+加羣目的。