在本文中，來自加拿大 Mila 研究所唐建課題組的研究人員提出了一種圖上高效能的嵌入訓練系統——GraphVite，訓練百萬級別的節點嵌入只需 1 分鐘左右，比現有實現快 50 倍以上。該系統最大可處理二十億邊的圖，是目前速度最快、規模最大的單機圖嵌入系統。
圖是一種通用、靈活的資料結構，可以用來編碼不同物件之間的關係，並且在現實世界中非常普遍，如社交網路、引文圖、蛋白質相互作用圖、知識圖譜等，涵蓋了多個應用和領域。最近，由於圖在多個任務中表現出色，越來越多的研究人員開始嘗試學習圖的有效表徵。然而，這是一個頗具挑戰性的問題，因為現實世界中的圖可能非常龐大，並且是異質的。因此，業界和學界的不同任務和應用都迫切需要可擴充套件的通用圖表徵系統。
為此，加拿大 Mila 研究所唐建課題組的研究人員開發了一個通用、高效能的圖嵌入系統——GraphVite。相較於 TensorFlow、PyTorch 等側重張量計算的機器學習系統，該系統針對大規模圖學習的特性，創新地將圖中的邊集劃分成多個部分，使多 GPU 能以非常小的同步代價進行訓練，並在多 GPU 上取得接近線性的加速比。架構中還將取樣部分移至 CPU 部分，既節省了超大圖的視訊記憶體開銷，又充分利用了 CPU 和主存的隨機訪問特性。
相比已有系統或實現，該系統將嵌入訓練速度提升了一至兩個數量級，只需約一分鐘的時間即可完成百萬級別的節點嵌入訓練，比現有實現快 50 倍以上。該系統對大圖上也有非常好的可擴充套件性，單機即可訓練十億級別的超大規模圖嵌入，是目前速度最快、規模最大的單機圖嵌入系統。目前，該專案已經開源。
論文：https://arxiv.org/pdf/1903.00757.pdf
網站：https://graphvite.io
GitHub 專案：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 等等。