Facebook 提出了一種可高效訓(xùn)練包含數(shù)十億節(jié)點(diǎn)和數(shù)萬億邊的圖模型的框架 BigGraph 并開源了其 PyTorch 實(shí)現(xiàn)。本文將解讀它的創(chuàng)新之處，解析它能從大規(guī)模圖網(wǎng)絡(luò)高效提取知識(shí)的原因。

圖（graph）是機(jī)器學(xué)習(xí)應(yīng)用中最基本的數(shù)據(jù)結(jié)構(gòu)之一。具體來說，圖嵌入方法是一種無監(jiān)督學(xué)習(xí)方法，可使用本地圖結(jié)構(gòu)來學(xué)習(xí)節(jié)點(diǎn)的表征。社交媒體預(yù)測(cè)、物聯(lián)網(wǎng)模式檢測(cè)或藥物序列建模等主流場(chǎng)景中的訓(xùn)練數(shù)據(jù)可以很自然地表征為圖結(jié)構(gòu)。其中每一種場(chǎng)景都可以輕松得到具有數(shù)十億相連節(jié)點(diǎn)的圖。圖結(jié)構(gòu)非常豐富且具有與生俱來的導(dǎo)向能力，因此非常適合機(jī)器學(xué)習(xí)模型。盡管如此，圖結(jié)構(gòu)卻非常復(fù)雜，難以進(jìn)行大規(guī)模擴(kuò)展應(yīng)用。也因此，現(xiàn)代深度學(xué)習(xí)框架對(duì)大規(guī)模圖數(shù)據(jù)結(jié)構(gòu)的支持仍非常有限。

Facebook 推出過一個(gè)框架 PyTorch BigGraph：https://github.com/facebookresearch/PyTorch-BigGraph，它能更快更輕松地為 PyTorch 模型中的超大圖結(jié)構(gòu)生成圖嵌入。

某種程度上講，圖結(jié)構(gòu)可視為有標(biāo)注訓(xùn)練數(shù)據(jù)集的一種替代，因?yàn)楣?jié)點(diǎn)之間的連接可用于推理特定的關(guān)系。這種方法遵照無監(jiān)督圖嵌入方法的模式，它可以學(xué)習(xí)圖中每個(gè)節(jié)點(diǎn)的向量表征，其具體做法是優(yōu)化節(jié)點(diǎn)對(duì)的嵌入，使得之間有邊相連的節(jié)點(diǎn)對(duì)的嵌入比無邊相連的節(jié)點(diǎn)對(duì)的嵌入更近。這類似于在文本上訓(xùn)練的 word2vec 的詞嵌入的工作方式。

當(dāng)應(yīng)用于大型圖結(jié)構(gòu)時(shí)，大多數(shù)圖嵌入方法的結(jié)果都相當(dāng)局限。舉個(gè)例子，如果一個(gè)模型有 20 億個(gè)節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)有 100 個(gè)嵌入?yún)?shù)（用浮點(diǎn)數(shù)表示），則光是存儲(chǔ)這些參數(shù)就需要 800 GB 內(nèi)存，因此很多標(biāo)準(zhǔn)方法都超過了典型商用服務(wù)器的內(nèi)存容量。這是深度學(xué)習(xí)模型面臨的一大挑戰(zhàn)，也是 Facebook 開發(fā) BigGraph 框架的原因。

PyTorch BigGraph

PyTorch BigGraph（PBG）的目標(biāo)是擴(kuò)展圖嵌入模型，使其有能力處理包含數(shù)十億節(jié)點(diǎn)和數(shù)萬億邊的圖。PBG 為什么有能力做到這一點(diǎn)？因?yàn)樗褂昧怂拇蠡緲?gòu)建模塊：

1. 圖分區(qū)，這讓模型不必完全載入到內(nèi)存中。
2. 在每臺(tái)機(jī)器上的多線程計(jì)算
3. 在多臺(tái)機(jī)器上的分布式執(zhí)行（可選），所有操作都在圖上不相連的部分進(jìn)行
4. 分批負(fù)采樣，當(dāng)每條邊 100 個(gè)負(fù)例時(shí)，可實(shí)現(xiàn)每臺(tái)機(jī)器每秒處理超過 100 萬條邊。

通過將圖結(jié)構(gòu)分區(qū)為隨機(jī)劃分的 P 個(gè)分區(qū)，使得可將兩個(gè)分區(qū)放入內(nèi)存中，PBG 解決了傳統(tǒng)圖嵌入方法的一些短板。舉個(gè)例子，如果一條邊的起點(diǎn)在分區(qū) p1，終點(diǎn)在分區(qū) p2，則它會(huì)被放入 bucket (p1, p2)。然后，在同一模型中，根據(jù)源節(jié)點(diǎn)和目標(biāo)節(jié)點(diǎn)將這些圖節(jié)點(diǎn)劃分到 P2 bucket。完成節(jié)點(diǎn)和邊的分區(qū)之后，可以每次在一個(gè) bucket 內(nèi)執(zhí)行訓(xùn)練。bucket (p1, p2) 的訓(xùn)練僅需要將分區(qū) p1 和 p2 的嵌入存儲(chǔ)到內(nèi)存中。PBG 結(jié)構(gòu)能保證 bucket 至少有一個(gè)之前已訓(xùn)練的嵌入分區(qū)。

PBG 的另一大創(chuàng)新是訓(xùn)練機(jī)制的并行化和分布式。PBG 使用 PyTorch 自帶的并行化機(jī)制實(shí)現(xiàn)了一種分布式訓(xùn)練模型，這用到了前面描述的模塊分區(qū)結(jié)構(gòu)。在這個(gè)模型中，各個(gè)機(jī)器會(huì)協(xié)調(diào)在不相交的 bucket 上進(jìn)行訓(xùn)練。這會(huì)用到一個(gè)鎖服務(wù)器（lock server），其負(fù)責(zé)將 bucket 分派給工作器（worker），從而盡可能地減少不同機(jī)器之間的通信。每臺(tái)機(jī)器都可以使用不同的 bucket 并行地訓(xùn)練模型。

在上圖中，機(jī)器 2 中的 Trainer 模塊向機(jī)器 1 上的鎖服務(wù)器請(qǐng)求了一個(gè) bucket，這會(huì)鎖定該 bucket 的分區(qū)。然后該 trainer 會(huì)保存它不再使用的所有分區(qū)并從共享分區(qū)服務(wù)器載入它需要的新分區(qū)，此時(shí)它可以將自己的舊分區(qū)釋放回鎖服務(wù)器。然后邊會(huì)從一個(gè)共享文件系統(tǒng)載入，并在沒有線程內(nèi)同步的情況下在多個(gè)線程上進(jìn)行訓(xùn)練。在一個(gè)單獨(dú)的線程中，僅有少量共享參數(shù)會(huì)與一個(gè)共享參數(shù)服務(wù)器持續(xù)同步。模型檢查點(diǎn)偶爾會(huì)從 trainer 寫入到共享文件系統(tǒng)中。這個(gè)模型允許使用至多 P/2 臺(tái)機(jī)器時(shí)，讓一組 P 個(gè) bucket 并行化。

PBG 一項(xiàng)不那么直接的創(chuàng)新是使用了分批負(fù)采樣技術(shù)。傳統(tǒng)的圖嵌入模型會(huì)沿真正例邊將隨機(jī)的「錯(cuò)誤」邊構(gòu)建成負(fù)訓(xùn)練樣本。這能顯著提升訓(xùn)練速度，因?yàn)閮H有一小部分權(quán)重必須使用每個(gè)新樣本進(jìn)行更新。但是，負(fù)例樣本最終會(huì)為圖的處理引入性能開銷，并最終會(huì)通過隨機(jī)的源或目標(biāo)節(jié)點(diǎn)「損害」真正的邊。PBG 引入了一種方法，即復(fù)用單批 N 個(gè)隨機(jī)節(jié)點(diǎn)以得到 N 個(gè)訓(xùn)練邊的受損負(fù)例樣本。相比于其它嵌入方法，這項(xiàng)技術(shù)讓我們能以很低的計(jì)算成本在每條邊對(duì)應(yīng)的許多負(fù)例上進(jìn)行訓(xùn)練。

要增加在大型圖上的內(nèi)存效率和計(jì)算資源，PBG 利用了單批 Bn 個(gè)采樣的源或目標(biāo)節(jié)點(diǎn)來構(gòu)建多個(gè)負(fù)例。在典型的設(shè)置中，PBG 會(huì)從訓(xùn)練集取一批 B=1000 個(gè)正例，然后將其分為 50 條邊一個(gè)的塊。來自每個(gè)塊的目標(biāo)（與源等效）嵌入會(huì)與從尾部實(shí)體類型均勻采樣的 50 個(gè)嵌入相連。50 個(gè)正例與 200 個(gè)采樣節(jié)點(diǎn)的外積等于 9900 個(gè)負(fù)例。

分批負(fù)采樣方法可直接影響模型的訓(xùn)練速度。如果沒有分批，訓(xùn)練的速度就與負(fù)例的數(shù)量成反比。分批訓(xùn)練可改善方程，得到穩(wěn)定的訓(xùn)練速度。

Facebook 使用 LiveJournal、Twitter 數(shù)據(jù)和 YouTube 用戶互動(dòng)數(shù)據(jù)等不同的數(shù)據(jù)集評(píng)估了 PBG。此外，PBG 還使用 Freebase 知識(shí)圖譜進(jìn)行了基準(zhǔn)測(cè)試，該知識(shí)圖譜包含超過 1.2 億個(gè)節(jié)點(diǎn)和 27 億條邊。另外還使用 Freebase 的一個(gè)小子集 FB15k 進(jìn)行了測(cè)試，F(xiàn)B15k 包含 15000 個(gè)節(jié)點(diǎn)和 600000 條邊，常被用作多關(guān)系嵌入方法的基準(zhǔn)。FB15k 實(shí)驗(yàn)表明 PBG 的表現(xiàn)與當(dāng)前最佳的圖嵌入模型相近。但是，當(dāng)在完整的 Freebase 數(shù)據(jù)集上評(píng)估時(shí)，PBG 的內(nèi)存消耗得到了 88% 的改善。

PBG 是首個(gè)可擴(kuò)展的、能訓(xùn)練和處理包含數(shù)十億節(jié)點(diǎn)和數(shù)萬億邊的圖數(shù)據(jù)的方法。PBG 的首個(gè)實(shí)現(xiàn)已經(jīng)開源，未來應(yīng)該還會(huì)有更有意思的貢獻(xiàn)。