訓練超大規模圖模型，PyTorchBigGraph如何作到？

Facebook 提出了一種可高效訓練包含數十億節點和數萬億邊的圖模型的框架 BigGraph 並開源了其 PyTorch 實現。本文將解讀它的創新之處，解析它能從大規模圖網絡高效提取知識的緣由。

 

圖（graph）是機器學習應用中最基本的數據結構之一。具體來講，圖嵌入方法是一種無監督學習方法，可以使用本地圖結構來學習節點的表徵。社交媒體預測、物聯網模式檢測或藥物序列建模等主流場景中的訓練數據能夠很天然地表徵爲圖結構。其中每一種場景均可以輕鬆獲得具備數十億相連節點的圖。圖結構很是豐富且具備與生俱來的導向能力，所以很是適合機器學習模型。儘管如此，圖結構卻很是複雜，難以進行大規模擴展應用。也所以，現代深度學習框架對大規模圖數據結構的支持仍很是有限。

Facebook 推出過一個框架 PyTorch BigGraph：https://github.com/facebookresearch/PyTorch-BigGraph，它能更快更輕鬆地爲 PyTorch 模型中的超大圖結構生成圖嵌入。

某種程度上講，圖結構可視爲有標註訓練數據集的一種替代，由於節點之間的鏈接可用於推理特定的關係。這種方法遵守無監督圖嵌入方法的模式，它能夠學習圖中每一個節點的向量表徵，其具體作法是優化節點對的嵌入，使得之間有邊相連的節點對的嵌入比無邊相連的節點對的嵌入更近。這相似於在文本上訓練的 word2vec 的詞嵌入的工做方式。

當應用於大型圖結構時，大多數圖嵌入方法的結果都至關侷限。舉個例子，若是一個模型有 20 億個節點，每一個節點有 100 個嵌入參數（用浮點數表示），則光是存儲這些參數就須要 800 GB 內存，所以不少標準方法都超過了典型商用服務器的內存容量。這是深度學習模型面臨的一大挑戰，也是 Facebook 開發 BigGraph 框架的緣由。

 

 

PyTorch BigGraph

PyTorch BigGraph（PBG）的目標是擴展圖嵌入模型，使其有能力處理包含數十億節點和數萬億邊的圖。PBG 爲何有能力作到這一點？由於它使用了四大基本構建模塊：

1. 圖分區，這讓模型沒必要徹底載入到內存中。
2. 在每臺機器上的多線程計算
3. 在多臺機器上的分佈式執行（可選），全部操做都在圖上不相連的部分進行
4. 分批負採樣，當每條邊 100 個負例時，可實現每臺機器每秒處理超過 100 萬條邊。

經過將圖結構分區爲隨機劃分的 P 個分區，使得可將兩個分區放入內存中，PBG 解決了傳統圖嵌入方法的一些短板。舉個例子，若是一條邊的起點在分區 p1，終點在分區 p2，則它會被放入 bucket (p1, p2)。而後，在同一模型中，根據源節點和目標節點將這些圖節點劃分到 P2 bucket。完成節點和邊的分區以後，能夠每次在一個 bucket 內執行訓練。bucket (p1, p2) 的訓練僅須要將分區 p1 和 p2 的嵌入存儲到內存中。PBG 結構能保證 bucket 至少有一個以前已訓練的嵌入分區。

PBG 的另外一大創新是訓練機制的並行化和分佈式。PBG 使用 PyTorch 自帶的並行化機制實現了一種分佈式訓練模型，這用到了前面描述的模塊分區結構。在這個模型中，各個機器會協調在不相交的 bucket 上進行訓練。這會用到一個鎖服務器（lock server），其負責將 bucket 分派給工做器（worker），從而儘量地減小不一樣機器之間的通訊。每臺機器均可以使用不一樣的 bucket 並行地訓練模型。

在上圖中，機器 2 中的 Trainer 模塊向機器 1 上的鎖服務器請求了一個 bucket，這會鎖定該 bucket 的分區。而後該 trainer 會保存它再也不使用的全部分區並從共享分區服務器載入它須要的新分區，此時它能夠將本身的舊分區釋放回鎖服務器。而後邊會從一個共享文件系統載入，並在沒有線程內同步的狀況下在多個線程上進行訓練。在一個單獨的線程中，僅有少許共享參數會與一個共享參數服務器持續同步。模型檢查點偶爾會從 trainer 寫入到共享文件系統中。這個模型容許使用至多 P/2 臺機器時，讓一組 P 個 bucket 並行化。

 

PBG 一項不那麼直接的創新是使用了分批負採樣技術。傳統的圖嵌入模型會沿真正例邊將隨機的「錯誤」邊構建成負訓練樣本。這能顯著提高訓練速度，由於僅有一小部分權重必須使用每一個新樣本進行更新。可是，負例樣本最終會爲圖的處理引入性能開銷，並最終會經過隨機的源或目標節點「損害」真正的邊。PBG 引入了一種方法，即複用單批 N 個隨機節點以獲得 N 個訓練邊的受損負例樣本。相比於其它嵌入方法，這項技術讓咱們能以很低的計算成本在每條邊對應的許多負例上進行訓練。

要增長在大型圖上的內存效率和計算資源，PBG 利用了單批 Bn 個採樣的源或目標節點來構建多個負例。在典型的設置中，PBG 會從訓練集取一批 B=1000 個正例，而後將其分爲 50 條邊一個的塊。來自每一個塊的目標（與源等效）嵌入會與從尾部實體類型均勻採樣的 50 個嵌入相連。50 個正例與 200 個採樣節點的外積等於 9900 個負例。

分批負採樣方法可直接影響模型的訓練速度。若是沒有分批，訓練的速度就與負例的數量成反比。分批訓練可改善方程，獲得穩定的訓練速度。

Facebook 使用 LiveJournal、Twitter 數據和 YouTube 用戶互動數據等不一樣的數據集評估了 PBG。此外，PBG 還使用 Freebase 知識圖譜進行了基準測試，該知識圖譜包含超過 1.2 億個節點和 27 億條邊。另外還使用 Freebase 的一個小子集 FB15k 進行了測試，FB15k 包含 15000 個節點和 600000 條邊，常被用做多關係嵌入方法的基準。FB15k 實驗代表 PBG 的表現與當前最佳的圖嵌入模型相近。可是，當在完整的 Freebase 數據集上評估時，PBG 的內存消耗獲得了 88% 的改善。

PBG 是首個可擴展的、能訓練和處理包含數十億節點和數萬億邊的圖數據的方法。PBG 的首個實現已經開源，將來應該還會有更有意思的貢獻。

 

 

【編輯推薦】

1. 智能物聯網：將人工智能的力量帶入物聯網
2. 解決數獨問題用人工智能仍是量子計算？
3. 人工智能的企業家來講，這四個新的方向可能值得重視
4. 機器學習中的標籤泄漏介紹及其如何影響模型性能
5. 將來，你是「透明的」！計算機經過人腦信號預測想法，還原人腦圖像準確率高達83％