可擴展的圖神經結構搜索系統
一、問題
1、圖數據
在現實生活中,很多的數據都是以圖的形式存在,像社交網絡,知識譜藥物和新材料等,圖神經網絡也被廣泛的應用于多個場景,如推薦系統,異常檢測,藥物以及蛋白質結構預測等。
首先我們來對最常見圖卷積神經網絡做一個簡單的回顧,從公式上來看,GCN 的表達形式與傳統的深度神經網絡區別在于多了含自環的度矩陣和鄰接矩陣,也就是增加了聚合鄰居節點特征的一個過程。因此在 GCN 的每一層包含兩個操作,propagation 操作用來聚合鄰居的信息,transformation 操作用來做變換。如果公式中的 A 矩陣是單位矩陣或者刪去圖里所有的邊,那么GCN 此時就退化成了 MLP。
2、圖神經網絡
GCN 的性能優于 MLP 主要是由于使用了消息傳播機制,聚合來自于高階領域的信息,從而增強了自身的表達能力。目前除了常見的圖卷積神經網絡之外,還有圖注意網絡(GAT, Graph Attention Network),引入了注意力機制,給每個鄰居節點分配不同權重,從而優化聚合鄰居節點的表達形式;GraphSAGE 使用歸納式學習,以 mini batch 進行訓練,能進一步擴展應用到較大的圖數據集。
3、Neural Message Passing(消息傳遞機制)
傳統的 GNN 都遵循消息傳遞機制的范式,而這種機制主要分為兩部分:通信+計算。當神經網絡去聚合鄰居節點的信息的時候,就涉及到了通信的問題。計算就是利用聚合到的信息去更新節點,這一部分在各種神經網絡中都是必不可少的,圖神經網絡也不例外。這里重點在于通信的過程。
在工業界的場景下,數據的規模都是比較龐大的,我們會采用分布式的方法去部署和訓練節點。當一個節點和它的鄰居節點沒有部署在同一臺機器上,做聚合信息的操作時就需要去拉取另一臺機器上的鄰居節點的信息,就這涉及到了通信開銷,特別是在帶寬不好的情況下,通信開銷將會很大。圖神經網絡每一層都需要做這種拉取,同時在大規模的圖數據上,每一個 epoch,即每一次訓練,也會存在很高的通信開銷。
4、GNN 系統
雖然消息傳播機制具有如上的弊端,但是目前大部分的 GNN 系統還是默認去使用這一機制,比如最常用的 DGL 和 PyG,以兼容更多的模型。
在工業界,大規模的圖數據處理會帶來一系列挑戰:
首先就是較高的建模門檻。當接到一個任務后,我們需要進行建模,即編寫對應的代碼,設計一個網絡結構。在這個過程中會涉及到手動調參,更新訓練過程及參數進行驗證。這個過程事實上有很高的門檻,同時也非常耗時,使用大規模的圖數據訓練神經網絡甚至會需要一兩天的時間,在沒有先驗知識的情況下,調參的過程會非常緩慢,因此針對任務設計 GNN 就需要知識豐富的專家。
另一個挑戰就是低可擴展性,正如前面提到,在分布式上部署會有很大的通信開銷,事實上在單機上的計算開銷也很大,這就導致了很高的模型訓練時間和預測時間。后續會有實驗結果進行進一步的說明。
5、瓶頸
為了更加直觀理解通信瓶頸,這里做了一個簡單的實驗。
受制于單機的存儲和分布式通訊開銷,現有的消息傳遞機制不能很好的傳播到大圖上。圖 a 是我們使用 2 層的 GraphSAGE 訓練 Reddit 上的一個數據集,橫坐標為機器數量,縱坐標是加速比,可以看到在增加更多機器時候,加速比增長并不明顯,甚至在機器從 2 個增加到 8 個時,加速比甚至沒有增加到 2 倍。圖 b 分析了這種情況的原因,即使增加機器后計算成本減小,更多的拉取數據操作使得通信開銷增大,因此加速比的增加并不明顯。
因此,我們的目標就是去如何兼容 GNN 的可擴展性,設計使用門檻低的圖神經網絡系統。
二、方法
1、系統目標
我們的 PaSca 自動搜索系統是一個端對端的設計,不需要取人為定義網絡結構和訓練流程,只需要輸入圖數據和優化目標,就可以輸出能兼顧多個優化目標的可擴展性 GNN,同。輸入的圖數據可以是特征標簽、圖結構矩陣,優化目標包括學術界常應用的分類性能指標,如準確率,AUC 等,同時還可以包括一些其他的約束,比如期望的訓練時長,期望使用的機器內存,從而達到控制預算的效果。也就是說,我們的系統可以考慮到與硬件相關的多個優化指標,在多個目標的約束下去完成多目標的搜索。
2、消息傳遞(Message Passing)范式
從數據流動的角度來看,GCN 遵從消息傳遞范式,它主要是從粒度較細的節點層次來刻畫數據的流動,主要由 3 個操作組成:message function 定義如何去生成信息,對應表達式中的度矩陣和鄰接矩陣;aggregation function 定義如何去聚合鄰居節點信息;update function定 義更新中心節點特征的方式,它遵從不斷地去迭代“聚合-更新”這一流程。
3、SGAP 建模范式(Scalable Graph Neural Architecture Paradigm)
(1)方法概覽
針對消息傳播機制的范式的弊端,我們提出了一種新的可擴展性的訓練流程的抽象,SGAP 建模范式(Scalable Graph Neural Architecture Paradigm)。不可擴展性的消息傳播機制范式需要在每個 epoch 訓練的時候拉取信息, 不斷地聚合-迭代鄰居的信息,它的通信次數是和訓練的 epoch 成正比的關系。而我們提出的 SGAP 建模范式主要分為三部分:預處理、模型訓練和后處理。預處理階段聚合鄰居的信息,聚合后我們完全拋棄了圖結構,將聚合后的信息傳入第二部分模型訓練中。這一步的模型可以是任意的機器學習模型,如 MLP,DNN,樹模型或更加傳統的神經網絡模型。得到模型預測之后再次調用一個聚合操作,從而得到最終的輸出。與之前不斷循環往復的過程相比,無論訓練的 epoch 是多少,SGAP 只做了兩次聚合鄰居的操作:預處理和后處理。
因此在分布式訓練的情況下,采用 SGAP 建模范式的 GNN 通信開銷會更低。對于分布式場景我們更加看重可擴展性,對于單機場景更看重效率。SGAP 在單機的情況下訓練也能顯著降低計算成本。具體來說,它在預處理階段減少了圖數據重復計算鄰接矩陣與特征矩陣相乘的步驟,避免了多次循環往復的聚合更新信息。矩陣的乘法運算在預處理階段計算好之后,再傳入模型訓練部分,按照普通的神經網絡模型計算即可,因此在單機模式下計算次數會變少,計算速度會提升。
(2)SGAP 范式
在 SGAP 范式中,后處理與預處理是比較類似的過程,下面將對 SGAP 的前兩個過程做更加數學化的闡述。它首先會在圖的層次做信息傳播,得到不同傳播層數的消息。接下來會聚合不同傳播層數的消息來的到新的特征,最后把得到的特征送入任意一個待訓練的機器學習模型。因此 SGAP 范式可以理解為從圖的層次去刻畫數據的流動,而不是更細粒度的像消息傳遞機制一樣從節點的層次。
(3)SGAP 抽象
對于預處理來說,從鄰居節點聚合的是特征,而后處理從鄰居節點聚合的是模型的輸出,也就是軟標簽。如圖中 A 點在第一次聚合時的到 BCD 三個鄰居節點的信息,然后再次調用通過 C 節點聚合到 E 和 F 節點的信息。同時已經被聚合的節點信息將被再次聚合。
(4)圖聚合器
在后處理這一步,我們參考了三種主流的圖聚合器,如 GCN 中使用的增強歸一化鄰接矩陣(Augmented normalized adjacency),APPNP 中使用的個性化 PageRank,和 MotifNet 中的 Triangle-induced adjacency。
(5)訓練
接下來訓練模型主要分為兩步,聚合來自預處理階段的消息,以及更新聚合后的信息。
(6)消息聚合器
消息聚合器分為兩類,第一種非自適應聚合器,比如在聚合的消息中取最大值、平均值等;第二種自適應的聚合器給不同節點的不同層表示消息不同的權重,如同注意力機制。熱力圖中顯示的是模型運行 50 次后的準確率,顏色越深代表準確性更高。
在圖中不同的節點所處的位置不同,權重也不盡相同,類似于社交網絡中活躍度和好友比較多的用戶,比如 18 號,只傳播 1 個 step 就達到了很高的準確率,代表著傳播一次就幾乎覆蓋到全圖;而 4 號代表著影響力較為一般的用戶,它的分布比較稀疏,就需要更多的傳播次數才能提升準確率。同時在自適應聚合器中引入的其他參數也是有代價的,會對計算效率產生一定的影響,這里就是在用計算效率去換更高的準確率。
(7)基于 SGAP 范式來設計 GNN
基于 SGAP 范式來設計的 GNN,完全拋棄了消息傳播機制的循環迭代過程,主要分為三步:前處理、訓練、后處理。前處理利用圖聚合器聚合鄰居節點的信息存入矩陣;訓練階段利用消息聚合器聚合預處理階段的特征矩陣,設定信息更新模型(如 MLP)來學習節點的軟標簽類別分布;后處理就是基于模型更新后的預測結果當作新的特征,再一次使用圖聚合器來聚合鄰居標簽信息,最后輸出最終預測。
4、自動化搜索系統(PaSca)
基于可擴展的 SGAP 范式,我們提出了自動化搜索系統 PaSca,其中包含兩個模塊,自動化的搜索引擎和分布式的評估引擎。搜索引擎的主要的目就是推薦一個 configuration instance,即一個網絡結構配置,可以配置出一個可擴展的圖神經網絡結構。評估引擎將會評估所推薦的 configuration instance,真正在數據集上訓練并產生預測,然后根據驗證集的結果去評估準確率。
(1)搜索引擎
搜索引擎需要接收輸入的優化目標,處理不同優化目標之間的 tradeoff。優化目標包括對于模型結構的要求(模型本身結構)、模型性能的要求(準確率等性能指標)和硬件設置(資源、內存空間、時間)相關的約束。搜索引擎可以在多個目標的約束下,去搜索并推薦最佳的滿足多個優化目標的網絡結構配置。
(2)設計空間
PaSca 的設計空間包含在預處理,模型訓練,以及后處理等 3 個階段的局部設計。拿模型訓練階段的設計空間為例,需要考慮的包括使用什么樣的聚合器(自適應/非自適應),使用多少層的 dens layer 等。在每個階段都有 2 個參數可供選擇,三個階段共有超過 150k 種可能的 configuration instances,同時現有的 Scalable GNN 都可以在我們的設計空間里找到。
(3)推薦服務器
對于推薦服務器來說,PaSca 使用經典的基于貝葉斯優化的 SMBO 方法。整個過程分為三步:建模,推薦和更新。首先進行建模,學習配置(網絡結構)與優化目標之間的關系;建模完成后,推薦引擎去推薦能兼顧多個優化目標的配置;之后評估引擎評估該配置的效果,把觀測到的歷史記錄返回推薦服務器中,進一步更新推薦模型,使得推薦越來越準確。
(4)評估引擎
對于分布式評估引擎來說,首先會使用圖數據聚合器對大圖做切分,切分好圖數據后再基于重復計算的方法來做數據聚合。具體來說就是根據已經計算好的第(i)步消息來計算第(i+1)步消息。網絡結構訓練器相對簡單,它使用 MLP 基于 Mini-batch 來訓練網絡,并且基于 parameter server,去異步更新網絡參數。
三、實驗
1、實驗設置
我們在目前廣泛使用的經典數據集及數據量稍大的數據集和自己的 Industry 數據集上都做了驗證,主要是為了驗證以下就是 3 個目標。首先,我們希望提出的 SGAP 比基于 NMP 的消息傳遞機制更具有擴展性;第二個目標是期望 PaSca 搜索出來的結果能夠很好地處理不同搜索目標之間的 tradeoff;第三個目標是希望 PaSca 搜索出來網絡結構能夠有更好的預測性能。
2、可擴展性分析
對于第一個目標,由于是針對大規模的數據去做設計,最關心的其實不是模型的性能,是而是 scalability,也即模型的可擴展性。如圖所示,我們使用了基于 SGAP的APPNP和基于 NMP的 GraphSAGE 在兩個不同的數據集上進行了對比,實驗發現基于 SGAP 的 GNN 可以取得接近線性的加速比,并且比基于 NMP 的消息傳遞機制更加接近理想的加速比。
3、搜索出來的代表性方法
基于 SGAP 范式搜索出來的代表性方法可以兼顧多個搜索目標之間的 tradeoff。如圖中帕累托平面所示,橫軸代表模型的準確率,縱軸代表預測時間。可以看到 PaSca-V3 取得了最低的預測誤差,代價就是帶來了與 PaSca-V2 相比更長的預測時間。目前表現 SOTA 的可擴展網絡結構 GBP,也可以在帕累托平面上被搜索出來。表 3 代表著搜索出來的 3 個不同版本的代表性網絡結構。給予網絡結構三大部分具體的參數,就可以確定唯一定義的網絡結構。
同時結果顯示,目前搜索出的代表性模型都能很好兼顧訓練時間與測試準確率。圖中藍框內代表的 PaSca 搜索出的模型,其中 PaSca-V2 和 PaSca-V3 都取得了較高的預測準確率,但是明顯需要更少的時間。即使拿性能最差的 SGC 與非 SGAP 模型相比,在準確率降低幅度不大的情況下,效率高出其他模型一個數量級以上。因此 SGAP 在工業范式界的大規模圖數據上表現是相當不俗的,尤其是對訓練時間要求比較高的情況下。
4、預測性能
在預測性能上,SGAP 和其他非擴展性的建模范式相比,具有相當競爭力的性能。基于 SGAP 范式實現的模型具有預處理、訓練、后處理這樣一種三段式的結構,與常見的 NMP 和 DNMP 范式下的模型相比,它們都取得了較好的預測性能。比如,使用 SGAP 范式下的 PaSca-V3 模型在不同的數據上進行測試都取得了最好的性能。因此,SGAP 范式在保證模型可擴展性的同時,不會降低模型的準確率。
四、總結
目前我們實現了能自動化建模 10 億節點的超大規模圖神經網絡系統,部署于騰訊太極機器學習平臺,并廣泛應用于視頻推薦和內容風控等場景,系統部分功能已在 Github 開源:?https://github.com/PKU-DAIR/SGL?,系統論文獲得 CCF A 類數據挖掘旗艦會議 WWW 2022 唯一“最佳學生論文獎”(中國第2個),相關工作刷新了國際圖學習榜單 OGB 的 3 項第一。
這里對我們的工作做一個總結。PaSca 作為一個新穎的構建和探索可擴展 GNNs 的網絡結構搜索系統,不僅僅研究單個的網絡結構設計。PaSca 搜索出來的代表性模型能夠在預測性能、效率以及可擴展性等多個方面超越現有的 SOTA GNN 模型。同時 PaSca 能夠幫助研究者來探索不同的 Scalable GNN 結構設計,并且理解不同設計的特點和功能。
?同時也對 PaSca 系統開源工作做一個介紹。前面主要討論的 SGAP 建模和網絡結構搜索這兩個部分的功能,已經集成到 SGL 系統工具包中,是 SGL 系統設計不可或缺的一部分。這也正是 SGL 系統的設計目標之一,即具有高可擴展性,能夠高效處理分布式還是單機場景下的超大規模工業圖數據。
第二個目標就是實現自動化,PaSca 系統可以對指定的多個優化目標去自動化的搜索推薦網絡結構。
第三個目標是在應用性上,可以針對常見的節點分類、聚類預測,后面也會支持倒藥物發現和推薦等場景實現一鍵調用,具有針對多個任務定制的用戶友好的接口。
另一個目標就是針對數據側的優化,關注包括噪聲處理、長尾分布不均衡處理、圖結構數據稀疏等一系列問題。同時,系統還會內置集成多種有效的提點方法。
以上就是本次講座的全部內容。
五、問答環節
Q1:如何保證 SGAP 這種方式的調整不會影響模型的效果,是否有理論上的支持?
A1:首先我們回答,前半部分。從實驗結果上可以看到,基本上不會影響效果,效果反而會更好;然后另外,從 Open Graph Benchmark 榜單也可以看到,現在大部分的數據集其實都可以歸納成 SGAP 里的一個分支,或者說歸納 SGAP 范式的前半部分,因為很多方法可能會缺少后處理的操作。所以說,SGAP 是不會去影響模型的效果。至于在理論上的支持,比如 SGC 及其他很多的工作,以及去年有一篇論文,是專門去研究比如說 Graph Augmented MLP 的一種方法,因此在理論上也是有一些支持的。如果對這個感興趣的話,我們后面也可以發一下相關的論文。
Q2:SGAP 范式是否能夠支持復雜的 GNN 模型?比如說在邊上有 n 計算的階梯模型。
A2:目前的話就是研究的工作還不多,是不能去直接去支持或者說接替,怎么去聚合鄰居的信息的算法。但是目前的處理辦法是做了一個折中,把不同跳的信息聚合時,相當于在做一個節點層次上的聚合。這樣一種方法在保證了可擴展性的同時,也利用了 Attention 的思想去學習節點不同層數的表示。但是目前沒有辦法做到像 GAT 那么細的粒度,去聚合節點跟節點之間的信息,這是很好的一個研究的 topic,可以在聚合階段利用一些啟發式的方法,或者說利用一些有理論保證的方法,去做一些帶注意力機制這樣一種思想的聚合算法。
Q3:PaSca 框架可以支持多個優化目標,并在一定約束下進行學習,這里是如何添加這種約束的呢?
A3:這個問題涉及貝葉斯優化,建模時會設計一個 x 和一個 y, x 就是各種網絡結構y的話可以是 accuracy。當我們去訓練一個模型,比如說像高斯過程、GBM、樹模型,或者隨機森林等,多目標就意味著一個x對應多個 y。實現方式是用我們實驗室自研的 OpenBox 工具包去做的,可以很方便的支持,多個約束目標的貝葉斯優化。
Q4:SGAP 這種方式的話,為什么說有更好的拓展性呢?
A4:更好的擴展性是最重要的一點,也是大家最關心的一點。我們在此再重新描述一下,前面提到的就是 NMP 消息傳遞機制這種范式,在每個 epoch 都需要去不斷迭代地去做聚合和更新的操作,這樣帶來兩個問題,第一個問題就是需要不斷的重復做這種稀疏矩陣乘 dense 矩陣的矩陣乘法,這就造成很高的計算成本,而另外由于不斷的聚合操作,會產生很高的通信成本,尤其是對于大規模圖數據。
但是在 SGAP 范式下,分成預處理、訓練以及后處理三個階段,所以這個時候通信成本就不受到 training epoch 大小影響。只需要在預處理和后處理階段做兩次通信;計算的話也是同理,矩陣乘法只需要在預處理和后處理階段去做。在模擬訓練的時候,會完全拋棄掉圖結構。這樣就兼顧了可擴展性和效率的問題,比 NMP 的消息傳遞機制會更好。