Facebook 人工智能團(tuán)隊(duì)已經(jīng)創(chuàng)建并正在開(kāi)放源代碼 PyTorch Biggraph

有效處理大規(guī)模圖對(duì)于促進(jìn)人工智能的研究和應(yīng)用至關(guān)重要，但特別是在工業(yè)應(yīng)用中的圖，包含數(shù)十億個(gè)節(jié)點(diǎn)和數(shù)萬(wàn)億個(gè)邊，這超出了現(xiàn)有嵌入系統(tǒng)的能力。

因此，F(xiàn)acebook 人工智能團(tuán)隊(duì)已經(jīng)創(chuàng)建并正在開(kāi)放源代碼 PyTorch Biggraph（PBG）。

PBG 是一個(gè)用于學(xué)習(xí)大規(guī)模圖嵌入的分布式系統(tǒng)，特別適用于處理具有多達(dá)數(shù)十億實(shí)體和數(shù)萬(wàn)億條邊的大型網(wǎng)絡(luò)交互圖。它在 2019 年的 SysML 會(huì)議上發(fā)表的大規(guī)模圖嵌入框架論文中提出。

PBG 比常用的嵌入軟件更快，并在標(biāo)準(zhǔn)基準(zhǔn)上生成與最先進(jìn)模型質(zhì)量相當(dāng)?shù)那度?。有了這個(gè)新工具，任何人都可以用一臺(tái)機(jī)器或多臺(tái)機(jī)器并行地讀取一個(gè)大圖并快速生成高質(zhì)量的嵌入。

PBG 對(duì)傳統(tǒng)的多關(guān)系嵌入系統(tǒng)進(jìn)行了多次修改，使其能夠擴(kuò)展到具有數(shù)十億個(gè)實(shí)體和數(shù)萬(wàn)億邊的圖。PBG 使用圖分區(qū)來(lái)在單個(gè)機(jī)器或分布式環(huán)境中訓(xùn)練任意量級(jí)的嵌入。研究人員在通用基準(zhǔn)測(cè)試中展示了與現(xiàn)有嵌入系統(tǒng)相當(dāng)?shù)男阅埽瑫r(shí)允許在多臺(tái)機(jī)器上擴(kuò)展到任意大的圖和并行化。他們?cè)趲讉€(gè)大型社會(huì)網(wǎng)絡(luò)圖以及完整的 Freebase 數(shù)據(jù)集上訓(xùn)練和評(píng)估嵌入，其中包含超過(guò) 1 億個(gè)實(shí)體和 20 億條邊。

具體而言，PBG 通過(guò)攝取圖的邊列表來(lái)訓(xùn)練輸入圖，每條邊由其源實(shí)體和目標(biāo)實(shí)體以及可能的關(guān)系類型進(jìn)行標(biāo)識(shí)。它為每個(gè)實(shí)體輸出一個(gè)特征向量(嵌入)，試圖將相鄰實(shí)體放置在向量空間中彼此靠近，同時(shí)將未連接的實(shí)體分開(kāi)。因此，具有相似鄰近分布的實(shí)體最終將位于附近位置。

可以使用在訓(xùn)練中學(xué)習(xí)的參數(shù)(如果有的話)，用不同的方法配置每種關(guān)系類型來(lái)計(jì)算這個(gè)“接近度得分（proximity score）”，這允許在多個(gè)關(guān)系類型之間共享相同的基礎(chǔ)實(shí)體嵌入。

其模型的通用性和可擴(kuò)展性使得 PBG 能夠從嵌入文獻(xiàn)的知識(shí)圖譜中訓(xùn)練出多種模型，包括 TransE、RESCAL、DistMult 和 ComplEx。

PBG 的設(shè)計(jì)考慮到了規(guī)?；⑼ㄟ^(guò)以下方式實(shí)現(xiàn)：

圖分區(qū)（graph partitioning），這樣模型就不必完全加載到內(nèi)存中；

每臺(tái)機(jī)器上的多線程計(jì)算；

跨多臺(tái)機(jī)器的分布式執(zhí)行（可選），所有機(jī)器同時(shí)在圖的不相交部分上運(yùn)行；

批量負(fù)采樣（batched negative sampling），允許處理的數(shù)據(jù)為> 100 萬(wàn)邊/秒/機(jī)器。

作為一個(gè)示例，F(xiàn)acebook 還發(fā)布了包含 5000 萬(wàn)維基百科概念的 Wikidata 圖的首次嵌入版本，該圖用于 AI 研究社區(qū)中使用的結(jié)構(gòu)化數(shù)據(jù)。這些嵌入是用 PBG 創(chuàng)建的，可以幫助其他研究人員在維基數(shù)據(jù)概念上執(zhí)行機(jī)器學(xué)習(xí)任務(wù)。

需要注意的是，PBG不適用于小規(guī)模圖上具有奇怪模型的模型探索，例如圖網(wǎng)絡(luò)、深度網(wǎng)絡(luò)等。

安裝步驟及更多信息，請(qǐng)參考 GitHub 相關(guān)介紹和 PyTorch-BigGraph 文檔：

https://github.com/facebookresearch/PyTorch-BigGraph

https://torchbiggraph.readthedocs.io/en/latest/

建立數(shù)十億個(gè)節(jié)點(diǎn)的嵌入圖

圖是表示多種數(shù)據(jù)類型的核心工具。它們可以用來(lái)對(duì)相關(guān)實(shí)體的網(wǎng)絡(luò)進(jìn)行編碼，例如關(guān)于世界的事實(shí)。例如，像 Freebase 這樣的知識(shí)庫(kù)具有不同的實(shí)體（如“Stan Lee”和“New York City”），作為描述它們之間關(guān)系的節(jié)點(diǎn)和邊（例如“出生于”）。

圖嵌入方法通過(guò)優(yōu)化目標(biāo)來(lái)學(xué)習(xí)圖中每個(gè)節(jié)點(diǎn)的向量表示，即具有邊的節(jié)點(diǎn)對(duì)的嵌入比沒(méi)有共享邊的節(jié)點(diǎn)對(duì)更接近，這類似于 word2vec 等詞嵌入在文本上的訓(xùn)練方式。

圖嵌入是一種無(wú)監(jiān)督學(xué)習(xí)，因?yàn)樗鼈冎皇褂脠D結(jié)構(gòu)學(xué)習(xí)節(jié)點(diǎn)的表示，而不使用基于任務(wù)的節(jié)點(diǎn)“標(biāo)簽”。與文本嵌入一樣，這些表示可用于各種下游任務(wù)。

超大規(guī)模圖形嵌入

當(dāng)前，超大規(guī)模圖形有數(shù)十億個(gè)節(jié)點(diǎn)和數(shù)萬(wàn)億條邊，而標(biāo)準(zhǔn)的圖嵌入方法不能很好地?cái)U(kuò)展到對(duì)超大規(guī)模圖的操作，這主要有兩大挑戰(zhàn)：首先，嵌入系統(tǒng)必須足夠快，以便進(jìn)行實(shí)際的研究和生產(chǎn)使用。例如，利用現(xiàn)有的方法，訓(xùn)練一個(gè)具有萬(wàn)億條邊的圖可能需要幾周甚至幾年的時(shí)間。

另外，存儲(chǔ)也是一大挑戰(zhàn)。例如，嵌入每個(gè)節(jié)點(diǎn)具有 128 個(gè)浮點(diǎn)參數(shù)的 20 億個(gè)節(jié)點(diǎn)，這需要 1TB 的數(shù)據(jù)，超過(guò)了商用服務(wù)器的內(nèi)存容量。

PBG 使用圖的塊分區(qū)來(lái)克服圖嵌入的內(nèi)存限制。節(jié)點(diǎn)被隨機(jī)劃分為 P 分區(qū)，這些分區(qū)的大小可以使內(nèi)存容納兩個(gè)分區(qū)。然后，根據(jù)邊的源節(jié)點(diǎn)和目標(biāo)節(jié)點(diǎn)，將邊劃分為 P2 簇（Buckets）。

大規(guī)模圖的 PBG 分區(qū)方案。節(jié)點(diǎn)被劃分為 P 分區(qū)，分區(qū)大小適合內(nèi)存。邊根據(jù)其源節(jié)點(diǎn)和目標(biāo)節(jié)點(diǎn)的分區(qū)劃分為簇。在分布式模式下，可以并行執(zhí)行具有非重疊分區(qū)的多個(gè)存儲(chǔ)簇（如藍(lán)色方塊所示）。

節(jié)點(diǎn)和邊進(jìn)行分區(qū)之后，就可以一次在一個(gè)簇上執(zhí)行訓(xùn)練。bucket（i,j）的訓(xùn)練只需要將節(jié)點(diǎn)分區(qū) i 和 j 嵌入存儲(chǔ)在內(nèi)存中。

PBG 提供了兩種方法來(lái)訓(xùn)練分區(qū)圖數(shù)據(jù)的嵌入。在單機(jī)訓(xùn)練中，嵌入件和邊在不使用時(shí)被交換到磁盤上。在分布式訓(xùn)練中，嵌入分布在多臺(tái)機(jī)器的內(nèi)存中。

分布式訓(xùn)練

PBG 使用 PyTorch 并行化原語(yǔ)（parallelization primitives）進(jìn)行分布式訓(xùn)練。由于一個(gè)模型分區(qū)一次只能由一臺(tái)機(jī)器調(diào)用，因此一次最多可以在 P/2 機(jī)器上訓(xùn)練嵌入。只有當(dāng)機(jī)器需要切換到新的簇時(shí)，模型數(shù)據(jù)才會(huì)進(jìn)行通信。對(duì)于分布式訓(xùn)練，我們使用經(jīng)典參數(shù)服務(wù)器模型，同步表示不同類型邊的共享參數(shù)。

APBG 分布式訓(xùn)練體系結(jié)構(gòu)。我們使用鎖服務(wù)器協(xié)調(diào)機(jī)器在不相交的存儲(chǔ)簇上進(jìn)行訓(xùn)練。分區(qū)模型參數(shù)通過(guò)分片分區(qū)服務(wù)器交換，共享參數(shù)通過(guò)分片參數(shù)服務(wù)器異步更新。

負(fù)采樣

圖嵌入和文本嵌入相似，構(gòu)造隨機(jī)的“假”邊與真正的邊一起作為負(fù)訓(xùn)練樣例。這大大加快了訓(xùn)練速度，因?yàn)槊總€(gè)新樣本只需更新一小部分權(quán)重。通常，這些消極的例子是由隨機(jī)源節(jié)點(diǎn)或目標(biāo)節(jié)點(diǎn)的“腐蝕”真邊構(gòu)成的。然而，我們發(fā)現(xiàn)對(duì)標(biāo)準(zhǔn)負(fù)抽樣的一些修改對(duì)于大規(guī)模圖是必要的。

首先，我們注意到在傳統(tǒng)的圖嵌入方法中，幾乎所有的訓(xùn)練時(shí)間都花在了負(fù)邊上。我們利用函數(shù)形式的線性特點(diǎn)，重用一批 N 個(gè)隨機(jī)節(jié)點(diǎn)，生成 N 個(gè)訓(xùn)練邊的損壞負(fù)樣本。與其他嵌入方法相比，此技術(shù)允許我們以很小的計(jì)算成本在每個(gè)真邊上訓(xùn)練許多負(fù)示例。

我們還發(fā)現(xiàn)，為了生成在各種下游任務(wù)中有用的嵌入，一種有效的方法是破壞邊，將 50% 的節(jié)點(diǎn)和另外 50% 的節(jié)點(diǎn)（根據(jù)其邊數(shù)進(jìn)行采樣）混合在一起。

最后，我們引入了“實(shí)體類型”的概念，它限制了如何使用節(jié)點(diǎn)構(gòu)造負(fù)樣本。例如，考慮一個(gè)包含歌曲、藝術(shù)家和流派節(jié)點(diǎn)的圖，并假設(shè)藝術(shù)家和歌曲之間存在“創(chuàng)作”關(guān)系。如果我們?yōu)檫@個(gè)關(guān)系統(tǒng)一抽樣源實(shí)體，我們將絕大多數(shù)抽樣歌曲（因?yàn)楦枨人囆g(shù)家多），但這些不是有效的潛在邊（因?yàn)楦枨荒苡伤囆g(shù)家制作）。PBG 可以基于關(guān)系的實(shí)體類型限制構(gòu)造哪些負(fù)樣本。

評(píng)估 PyTorch-BigGraph

為了評(píng)估 PBG 的性能，我們使用了公開(kāi)的 Freebase 知識(shí)圖，它包含超過(guò) 1.2 億個(gè)節(jié)點(diǎn)和 27 億條邊。我們還使用了一個(gè)較小的 Freebase 圖子集（FB15K），它包含 15000 個(gè)節(jié)點(diǎn)和 600000 條邊，通常用作多關(guān)系嵌入方法的基準(zhǔn)。

T-SNE 繪制的由 PBG 訓(xùn)練的 Freebase 知識(shí)圖嵌入。國(guó)家、數(shù)字和科學(xué)期刊等實(shí)體也有類似的嵌入。

可以看出，對(duì)于 FB15k 數(shù)據(jù)集，PBG 和最新的嵌入方法性能相當(dāng)。

圖：FB15K 數(shù)據(jù)集的鏈路預(yù)測(cè)任務(wù)上嵌入方法的性能。PBG 使用其模型來(lái)匹配 transe 和復(fù)雜嵌入方法的性能。我們測(cè)量了 MRR，并在 FB15K 測(cè)試集上對(duì)鏈接預(yù)測(cè)進(jìn)行 hit@10統(tǒng)計(jì)。Lacroix 等人使用非常大的嵌入維數(shù)實(shí)現(xiàn)更高的 MRR，我們可以在 PBG 中采用同樣的方法，但這里暫不涉及。

下面，我們使用 PBG 對(duì)完整的 Freebase 圖訓(xùn)練嵌入?，F(xiàn)代服務(wù)器可以容納這個(gè)規(guī)模的數(shù)據(jù)集 但 PGB 分區(qū)和分布式執(zhí)行既節(jié)約了內(nèi)存，也縮短了訓(xùn)練時(shí)間。我們發(fā)布了 Wikidata 的首次嵌入，這是一個(gè)相似數(shù)據(jù)中更新的知識(shí)圖。

我們還評(píng)估了幾個(gè)公開(kāi)的社交圖數(shù)據(jù)集的 PBG 嵌入，發(fā)現(xiàn) PBG 優(yōu)于其他競(jìng)爭(zhēng)方法，并且分區(qū)和分布式執(zhí)行減少了內(nèi)存使用和培訓(xùn)時(shí)間。對(duì)于知識(shí)圖、分區(qū)或分布式執(zhí)行使得訓(xùn)練對(duì)超參數(shù)和建模選擇更加敏感。然而對(duì)于社交圖來(lái)說(shuō)，嵌入質(zhì)量似乎對(duì)分區(qū)和并行化選擇并不敏感。

利用分布式訓(xùn)練的優(yōu)勢(shì)進(jìn)行嵌入

PBG 允許 AI 社區(qū)為大規(guī)模圖（包括知識(shí)圖譜）以及其他如股票交易圖、在線內(nèi)容圖和生物數(shù)據(jù)圖訓(xùn)練嵌入，而無(wú)需專門的計(jì)算資源（如 GPU 或大量?jī)?nèi)存）。我們還希望 PBG 將成為小型公司和機(jī)構(gòu)的有用工具，他們可能擁有大型圖數(shù)據(jù)集，但沒(méi)有將這些數(shù)據(jù)應(yīng)用到其 ML 應(yīng)用程序的工具。

雖然我們?cè)?Freebase 等數(shù)據(jù)集上演示了 PBG，但 PBG 真正的設(shè)計(jì)意圖是處理比此圖大 10~100 倍的圖。我們希望這能鼓勵(lì)實(shí)踐者發(fā)布和試驗(yàn)更大的數(shù)據(jù)集。計(jì)算機(jī)視覺(jué)（通過(guò)對(duì)標(biāo)簽的 Deep Learning 來(lái)改進(jìn)圖像識(shí)別質(zhì)量）和自然語(yǔ)言處理（word2vec、BERT、Elmo）的最新突破是對(duì)海量數(shù)據(jù)集進(jìn)行未知任務(wù)預(yù)訓(xùn)練的結(jié)果。我們希望通過(guò)對(duì)大規(guī)模圖的無(wú)監(jiān)督學(xué)習(xí)，最終能夠得到更好的圖結(jié)構(gòu)化數(shù)據(jù)推理算法。