知識圖譜 （Knowledge Graph）作為一個重要的技術，在近幾年裡被廣泛運用在了信息檢索，自然語言處理，以及推薦系統等各種領域。學習知識圖譜的嵌入表示 （Knowledge Graph Embeddings）是一種從知識圖譜結構中生成無監督節點特徵（node feature）的方法，生成的特徵可以被用在各種機器學習任務之上。例如，可以通過節點的嵌入表示來預測兩個節點之間是否有連結（link prediction）。
然而，隨著社交網絡、推薦系統等典型圖數據場景的發展，知識圖譜的規模也在不斷地增長。在工業界真實的場景中，技術人員常常需要面對千萬級，甚至是億萬級節點的大規模圖數據。如何快速、高效地在大規模知識圖譜上進行嵌入表示的訓練是當前的一個挑戰。
近日，亞馬遜 AI 團隊繼 DGL 之後，又開源了一款專門針對大規模知識圖譜嵌入表示的新訓練框架 DGL-KE，旨在能讓研究人員和工業界用戶方便、快速地在大規模知識圖譜數據集上進行機器學習訓練任務。
github 地址：https://github.com/awslabs/dgl-ke
相比於已有的開源框架，DGL-KE 的亮點如下：
支持各種主流知識圖譜表示學習算法，包括 TransE、ComplEx、DistMult、TransR、RESCAL、RotatE 等；已有開源框架中唯一支持多核 CPU，多卡 GPU，CPU-GPU 混合訓練，以及分布式訓練的知識圖譜嵌入表示框架；簡單易用，用戶無需編寫代碼，直接將知識圖譜數據作為輸入即可；高性能且可擴展。根據 DGL-KE 發布的 Freebase 數據集 （超過 8600 萬節點，3 億條邊）Benchmark 顯示，在 AWS EC2 平臺上，一臺 p3.16xlarge （8 GPUs）可以在 100 分鐘內完成訓練。4 臺 r5dn.24xlarge （4*48 CPUs）可以在 30 分鐘內完成訓練，並且達到線性加速比。這一結果比當前最快的同類系統（如 Facebook 發布 Pytorch-BigGraph）快 2-5 倍。
圖1: DGL-KE 系統架構
DGL-KE 之所以能夠有這樣的性能，主要是因為採用了許多創新的系統和算法優化：
（一）基於 METIS 圖分割算法的分布式訓練
對超大規模的圖數據進行訓練，分布式訓練必不可少。其思路主要是將原始大圖分割成不同的子圖，每一臺機器負責在一個子圖上進行隨機梯度下降訓練，所有機器之間通過參數伺服器（Parameter Server）進行模型的同步。其架構如下圖所示：
圖2: DGL 分布式架構
然而，如果只是對一張大圖進行隨機切割，會造成訓練機器與參數伺服器之間的數據通信量巨大（本地機器需要從遠程機器去請求自己所需的模型數據），從而造成網絡瓶頸。為了解決這一問題，DGL-KE 在訓練前會預先通過 METIS 圖分割算法對原始數據進行切割。
METIS 算法是計算機科學家 George Karypis 於 1995 年提出的一種高效的圖分割算法，而 George Karypis 也正是 DGL-KE 項目的作者之一。METIS 算法會將一張大圖上相關聯的節點儘可能放置在同一個劃分（partition）之中。這樣可以將大部份的網絡通信開銷轉化成本地機器的內存拷貝，從而大大提升了分布式訓練的速度。
在 Freebase 這個數據集的實際訓練中，METIS 算法可以節省將近 90% 的模型網絡傳輸帶寬，從而使分布式訓練達到線性加速比。DGL-KE 的分布式訓練使用了 DGL-KVStore 組件。DGL-KVStore 是專門為 DGL 系統定製開發的參數伺服器模塊，用來實現稀疏模型的同步通訊。該組件通過 C++ 實現了底層 socket、消息隊列，以及稀疏數據序列化的定向優化，並且可以無縫兼容 METIS 圖分割算法。
（二）基於共享內存的單機多進程訓練
多核（Multi-core）已經成為當前計算機體系架構的標配，很多強大的 workstation 在一臺機器內更是會有超過幾十個 CPU 核心和上百 GB 甚至上 T 的內存。對於很多千萬級節點的圖數據來說，這樣的單機性能已經足夠處理這種規模的數據。
DGL-KE 針對這樣的場景也做了相應的系統優化，讓用戶儘可能地挖掘一臺機器的性能極限。與傳統基於多線程（Multi-thread）的並行優化不同，DGL-KE 採用了基於多進程（Multi-Process）的粗粒度並行優化。粗粒度的並行可以最大限的提升程序運行並行度，從而提高加速比。此外，DGL-KE 在不同進程之間通過共享內存（Shared-memory）進行模型同步，從而大大減小了進程之間的通信開銷。
圖3: 基於共享內存的單機多進程訓練
（三）CPU-GPU 混合訓練
知識圖譜嵌入表示的訓練過程中會產生大量的矩陣運算，而矩陣運算可以通過 GPU 來加速。對於小規模的圖數據來說，DGL-KE 允許用戶將完整的圖模型放入 GPU 中進行訓練，從而達到最優性能。然而，相比於 CPU 內存，GPU 內存要小很多，一但模型 Embeddings 的大小超過了 GPU 內存限制就無法進行訓練。針對這樣的場景，DGL-KE 為用戶提供了 CPU-GPU 混合訓練的模式。
在 CPU-GPU 混合訓練模式中，模型 Embeddings 被存儲在 CPU 的內存裡，而 GPU 通過 mini-batch 的方式在每一輪迭代過程中將一小部分數據從 CPU 拷貝到 GPU 進行訓練。為了避免 CPU 與 GPU 之間的數據拷貝開銷，DGL-KE 採用異步訓練的方式將數據拷貝與計算 overlap 起來。然而，異步計算會帶來模型收斂速度和精確度的下降，DGL-KE 在這裡採用了另一個優化，將 Entity Embedding 和 Relation Embedding 的更新採用不同的方式進行：relation 使用同步更新，而 enity 使用異步更新。
之所以這樣做是因為在實際的訓練過程中，relation 在很多數據集上都表現為長尾分布，即某幾種 relation 類型佔據絕大多數，所以異步更新會導致 relation embedding 在訓練過程中產生大量的模型衝突，從而影響模型的收斂和準確性。而 entity 在訓練過程通常是稀疏的，所以異步訓練只會產生很小的衝突。採用這樣一個簡單的優化，DGL-KE 既可以保證模型訓練的收斂性，又可以保證系統性能。
圖4: CPU-GPU 混合訓練
除了以上優化之外，DGL-KE 還提供了其他若干優化方法。例如，使用 Joint Negative Sampler 加速負採樣過程，使用 Relation Partition 來減少訓練過程中的數據拷貝，以及使用 Periodic synchronization 保證模型的收斂等。DGL-KE 內置了多個處理好格式的知識圖譜數據集，用戶可以直接下載使用。
除此之外，DGL-KE 在兩個小數據集 FB15k，wn18，以及一個大數據集 Freebase 上提供了訓練 Benchmark，用戶可以直接通過提供的腳本復現訓練結果。和已有的開源框架相比，DGL-KE 的性能優勢明顯，如下結果展示了 DGL-KE 與 Graphvite 在 FB15k 數據集上的性能比較，以及和 Pytorch-Biggraph 在 Freebase 數據集上的比較。
DGL-KE vs Graphvite
DGL-KE vs Pytorch-Biggraph
【END】