機器之心發布
作者：Zhaocheng Zhu等
在本文中，來自加拿大 Mila 研究所唐建課題組的研究人員提出了一種圖上高性能的嵌入訓練系統——GraphVite，訓練百萬級別的節點嵌入只需 1 分鐘左右，比現有實現快 50 倍以上。該系統最大可處理二十億邊的圖，是目前速度最快、規模最大的單機圖嵌入系統。
圖是一種通用、靈活的數據結構，可以用來編碼不同對象之間的關係，並且在現實世界中非常普遍，如社交網絡、引文圖、蛋白質相互作用圖、知識圖譜等，涵蓋了多個應用和領域。最近，由於圖在多個任務中表現出色，越來越多的研究人員開始嘗試學習圖的有效表徵。然而，這是一個頗具挑戰性的問題，因為現實世界中的圖可能非常龐大，並且是異質的。因此，業界和學界的不同任務和應用都迫切需要可擴展的通用圖表徵系統。
為此，加拿大 Mila 研究所唐建課題組的研究人員開發了一個通用、高性能的圖嵌入系統——GraphVite。相較於 TensorFlow、PyTorch 等側重張量計算的機器學習系統，該系統針對大規模圖學習的特性，創新地將圖中的邊集劃分成多個部分，使多 GPU 能以非常小的同步代價進行訓練，並在多 GPU 上取得接近線性的加速比。架構中還將採樣部分移至 CPU 部分，既節省了超大圖的顯存開銷，又充分利用了 CPU 和主存的隨機訪問特性。
相比已有系統或實現，該系統將嵌入訓練速度提升了一至兩個數量級，只需約一分鐘的時間即可完成百萬級別的節點嵌入訓練，比現有實現快 50 倍以上。該系統對大圖上也有非常好的可擴展性，單機即可訓練十億級別的超大規模圖嵌入，是目前速度最快、規模最大的單機圖嵌入系統。目前，該項目已經開源。
論文：https://arxiv.org/pdf/1903.00757.pdf網站：https://graphvite.ioGitHub 項目：https://github.com/DeepGraphLearning/graphvite
GraphVite 的開源版本包含節點嵌入、知識圖譜嵌入和圖結構&高維可視化三個應用的實現：
節點嵌入。目的是學習大規模圖的節點表徵。GraphVite 目前包含一些 SOTA 節點嵌入方法，如 DeepWalk、LINE 和 node2vec，未來還會繼續增加。知識圖譜嵌入。目的是學習實體和關係的表徵，目前支持多種表徵方法，包括 TransE、DisMult、CompEx、SimplE 和 RotatE，之後還會加入更多方法。圖&高維數據可視化。GraphVite 還支持學習節點的二維或三維坐標來實現圖的可視化，還可以泛化到任何高維數據的可視化。這對於可視化深度神經網絡學到的表徵形式特別有用。GraphVite 目前已經實現了最先進的可視化算法之一——LargeVis，未來還會實現更多的可視化算法。
GraphVite 支持的三種應用。
除了驚人的速度，GraphVite 還為研發人員提供了用戶友好的完整應用 pipeline。該系統具有數據集和評估任務等模塊，是一個用於嵌入模型和實驗的自成體系的環境。其標準數據集中包含 30 多個現有模型的基線基準。只需一句指令即可復現這些基準，部署到更大的數據集也很容易，方便為圖表徵學習開發新的模型。
該研究顯著加快了圖嵌入、知識圖譜嵌入和圖結構&高維可視化的訓練速度，將多個標準數據集上的嵌入訓練時間刷新到了 15 分鐘左右，直接促進了圖上嵌入算法的實現與迭代，間接影響了圖表徵學習算法研發的範式。單 GPU 17 分鐘生成完整 ImageNet 數據集特徵空間的可視化，極大地方便了深度學習模型的調試和可視化。相較於同期 FAIR 的 PyTorchBigGraph 系統，單機多卡的設計更有助於學界的普及使用。該系統所支持的超大規模圖訓練將會為工業應用帶來不少機遇。
GraphVite 為什麼那麼快
GraphVite 根據 CPU 和 GPU 各自體系結構的特點，將圖嵌入訓練分為採樣和訓練兩個部分，分別交由 CPU 和 GPU 完成。其中採樣部分使用 CPU 並行在線增強，解決了現有算法中增強後的圖佔用內存過大的問題。訓練部分，系統提出了一種並行的負採樣方法。該方法將頂點的嵌入劃分為若干塊，並將每條邊按其兩端頂點進行分塊。在訓練過程中，多塊 GPU 始終在頂點集不相交的塊上工作。這一設計極大地減小了多 GPU 之間的同步代價，並使參數矩陣超出顯存的大規模嵌入訓練成為可能。此外，GraphVite 還採用協同策略，令多 CPU 和多 GPU 異步工作，從而使訓練速度翻倍。
並行在線增強
節點嵌入方法的第一階段是利用隨機遊動（random walk）來增強原始網絡。由於增強後的網絡通常要比原始網絡大一至兩個量級，因此若原始網絡已經足夠大，則不可能將增強後的網絡導入主存。因此，研究人員提出了並行在線增強，它能夠在線生成增強後的邊樣本，同時又無需顯示地存儲增強後的網絡。
這種方法可視為 LINE 中所用到的增強和邊採樣方法的一種在線擴展。首先，研究人員隨機採樣一個初始節點（depature node），其選擇概率與節點度成正比。然後，他們從初始節點執行隨機遊動操作，同時選擇特定增強距離內的節點對作為邊樣本。這裡要注意的是，同一隨機遊動中生成的邊樣本具有相關性，可能會使訓練效果變差。
受啟發於強化學習中廣泛使用的經驗回放方法，研究人員收集邊樣本到樣本池，並且在轉移樣本池至 GPU 進行嵌入訓練之前對它展開洗牌（shuffle）。當每個線程提前分配至單獨的樣本池時，文中提出的邊採樣方法可以進行並行化處理。
算法 2 展示了並行在線增強的具體過程。
儘管樣本池的洗牌對優化非常重要，但同時也減慢了網絡增強階段的運行速度（參見表 7）。原因在於：一般的洗牌包含大量隨機存儲訪問，無法通過 CPU 高速緩存獲得加速。如果伺服器不只有一個物理 CPU 插座（socket），則速度損失會更加嚴重。為了緩解這個問題，研究人員提出了偽洗牌（pseudo shuffle）方法，該方法以一種更有利於高速緩存的方式對關聯樣本進行洗牌，並顯著提升了系統運行速度。
並行負採樣
在嵌入訓練階段，研究者將訓練任務分解成小的片段，並將它們分配給多個 GPU。子任務的設計必須使用少量共享數據，以最小化 GPU 之間的同步成本。為了弄清如何將模型參數無重疊地分配到多個 GPU，研究者首先引入了一個-gradient exchangeable 的定義。
基於節點嵌入中觀察到的梯度可交換性，研究者提出了一個用於嵌入訓練階段的並行負採樣算法。對於 n 個 GPU，他們將頂點行和語境行分別劃分為 n 個分塊（見圖 2 左上角）。結果得到一個 n × n 的樣本池分塊網格，其中每條邊都屬於其中一個塊。如此一來，只要對每一塊施加迭代數量限制，任何一對不共享行或列的塊都是 gradient exchangeable。相同行或列裡的塊是-gradient exchangeable。
研究者將 episode 定義為並行負採樣中使用的塊級步驟。在每個 episode 中，他們分別向 n 個 GPU 發送 n 個不相交的塊及其對應的定點和語境分塊。接下來，每個 GPU 藉助 ASGD 更新自身的嵌入塊。由於這些塊是梯度可互換的，並且在參數矩陣中不共享任何行，因此多個 GPU 可以在不同步的情況下同時執行 ASGD。
在每個 episode 結尾，研究者從所有 GPU 中收集更新的權重並分配另外 n 個不相交的塊。此處的-gradient exchangeable 由 n 個不相交塊中的樣本總數來控制，研究者將其定義為 eposode size。eposode size 越小，嵌入訓練中的-gradient exchangeable 越好。
然而，較小的 eposode size 也會導致更頻繁的同步。因此，episode size 的大小需要在速度和-gradient exchangeable 之間進行權衡。下圖 2 給出了 4 個分塊的並行負採樣示例。
圖 2：在 4 個 GPU 上的並行負採樣示例。在每個 episode 期間，GPU 從樣本池中獲取正交塊。每個 GPU 利用從自身上下文節點中獲取的負樣本訓練嵌入。同步只需要在 episode 之間進行。
通常來講，節點嵌入方法從所有可能的節點中採樣負邊。然而，如果 GPU 必須通過互相通信才能獲得它們自己的負樣本，那將非常耗費時間。為了避免這一開銷，研究者對負樣本進行了限制，規定其只能從當前 GPU 語境行獲取。儘管這看起來有點問題，但實際效果非常好。
注意，雖然本文中列出的分區數等於 n，但並行負採樣可以很容易地推廣到分區數大於 n 的情況，只需在每個 episode 期間處理 n 的子組中的不相交塊即可。下圖中的算法 3 展示了這一多 GPU 的混合系統。
GPU 和 CPU 的協作策略
上述並行負採樣使得不同的 GPU 可以同時訓練節點嵌入，只需要在 episode 之間進行同步。然而，應該指出的是，GPU 和 CPU 之間也共享樣本池。如果它們在樣本池上同步，那麼只有同一階段的 worker 才能同時訪問樣本池，也就意味著硬體只在一半的時間處於理想狀態。
為了解決這一問題，研究者提出了一種協作策略來減少同步開銷。他們在主內存中分配了兩個樣本池，讓 CPU 和 GPU 始終在不同的樣本池上工作。CPU 首先填充一個樣本池並將其傳送給 GPU。然後，分別在 CPU 和 GPU 上並發執行並行在線增強和並行負採樣。當 CPU 填滿一個新池時，這兩個採樣池進行調換。下圖 1 展示了這一步驟。利用這種協作策略可以降低 CPU 和 GPU 之間的同步成本，同時將系統的速度加倍。
圖 1：本文中的混合系統概覽。灰色和黃色框分別對應網絡增強和嵌入訓練階段。這兩個階段藉助本文提出的協作策略異步執行。
實驗
研究人員利用實驗驗證了 GraphVite 的有效性和效率。他們首先在 Youtube 數據集（節點嵌入文獻中廣泛使用的大型網絡）上對系統進行評估，然後又在三個更大的數據集（Friendster-small、Hyperlink-PLD 和 Friendster）上評估了 GraphVite。
表 3：Youtube 數據集上不同系統的時間結果。預處理時間指的是訓練前的所有開銷，包括網絡輸入和離線網絡增強。需要注意的是，OpenNE 的預處理時間沒有可比性，因為它缺少了網路增強階段。GraphVite 的加速比基於當前最快的系統 LINE 進行計算。
表 4：Youtube 數據集上的節點分類結果。
表 5：更大數據集上的時間結果。研究人員只用 4 個 GPU 評估了 Hyperlink-PLD 和 Friendster 數據集，因為它們的嵌入矩陣無法讀入單個 GPU 的內存。
圖 4：GraphVite 在更大數據集上的性能曲線。對於 Friendster 數據集，研究人員繪製了 LINE 的結果以供參考。其他系統無法在一周內解決以上任何一個數據集。
控制變量研究
為了更全面地了解 GraphVite 的不同組件，研究人員進行了以下幾項控制變量實驗。為了呈現出直觀的比較，他們在標準 Youtube 數據集上對這些實驗進行評估。由於空間限制，研究人員只提供了 2% 標記數據基礎上的性能表現結果。所有的加速比根據 LINE 進行計算。
表 6：GraphVite 主要組件的控制變量。需要注意的是，基線具有與 GraphVite 相同的 GPU 實現效果，並在 CPU 上展開並行邊採樣。
表 7：不同洗牌算法下的性能表現結果和訓練時間。文中提出的偽洗牌（pseudo shuffle）算法在性能和速度之間達到了最佳權衡。
圖 5：不同 episode 大小下 GraphVite 的速度和性能表現。虛線表示未進行並行負採樣的單個 GPU 基線。
圖 6：不同硬體數量下的加速結果。加速幾乎與 CPU 和 GPU 的數量呈線性相關。
表 8：Tesla P100 伺服器和 GTX 1080 伺服器兩種硬體配置下的 GraphVite 訓練時間。
作者信息
這項工作由加拿大 Mila 研究所唐建老師小組的博士一年級學生 Zhaocheng Zhu 擔任一作，Shizhen Xu 和 Meng Qu 也參與了研究。
加拿大 Mila 研究所是由圖靈獎得主 Yoshua Bengio 帶領的深度學習研究所。作為學界最大的機器學習研究機構之一，Mila 專注於深度學習相關的基礎研究，以及學習算法在不同領域的應用。研究所在機器翻譯、物體識別和生成模型等方面有諸多為人熟知的工作，著名的《Deep Learning》一書亦出自 Mila。
唐建老師帶領的研究組是 Mila 活躍的圖學習研究小組，其主要方向包括圖表徵學習和深度學習的基礎研究，涵蓋知識圖譜、藥物發現和推薦系統等不同應用領域。研究組在圖表徵學習領域具有相當的知名度，代表作有 LINE，LargeVis 和 RotatE 等等。