訓練超大規模圖模型，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 為什么有能力做到這一點？因為它使用了四大基本構建模塊：
- 圖分區，這讓模型不必完全載入到內存中。
- 在每臺機器上的多線程計算
- 在多臺機器上的分布式執行（可選），所有操作都在圖上不相連的部分進行
- 分批負采樣，當每條邊 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 的首個實現已經開源，未來應該還會有更有意思的貢獻。
