PyTorch-BigGraph: A Large Scale Graph Embedding System Software Engineer ynqa

PyTorch-BigGraph - Billion-Trillionサイズのエッジを持つグラフに対して Graph Embeddingsを行う - For Knowledge Graph - シングルマシン、分散学習における省メモリ /計算効率の向上 - Negative Samplingにおける効率的な負例の生成 - OSS - https://github.com/facebookresearch/PyTorch-BigGraph

1. Graph Embeddings 2. PyTorch-BigGraph a. Data Structure b. Distributed Training c. Single Machine Training d. Negative Sampling e. Evaluation Outline

Graph Embeddings

- Graph Embeddingsといっても、実際には様々な Input/Outputを扱うモデルが存在する - Input: - Homogeneous Graph (e.g. undirected and unweighted graph) - Heterogeneous graph (e.g. Knowledge Graph) - Output: - Node, Edge, Whole-Graph - Methods - Random walk: DeepWalk, node2vec - Graph convolutional network - For Knowledge Graph: TransE, RESCAL, DistMult, ComplEx Graph Embeddings

- Embeddings - Word Embeddings: word -> vector - Graph Embeddings: node -> vector - Approach - グラフの各ノード（エンティティ）におけるベクトル表現を獲得する。あるエッジで結ばれた エンティティ同士は結ばれていないエンティティよりも似たようなベクトルを獲得するよう に学習する - Word EmbeddingsにおけるDistributional Hypothesisのようなもの Graph Embeddings

- のtripletで構成される有向グラフについて Embeddingを行う - 学習モデル: TransE, RESCAL, DistMult, ComplEx ⇒ もう少し一般化されて実装されている Knowledge Graph

Scoring : 複素数の実部（real part） : 複素共役（complex conjugate） : element-wise production

- Pairwise ranking loss (Triplet loss） - 正例とrelationの距離が近くなり、負例と relationの距離が遠くなるように学習する Objective functions

PyTorch BigGraph

Data Structure

- 1つのエンティティが1つのエンティティタイプを持ち、他のエンティ ティとのエッジを複数持ち、各エッジは 1つのリレーションを持つ - あるリレーションにおける左側 /右側のエンティティタイプはす べて同じでなければならない Data Model

- エンティティタイプ毎にさらにパーティションに区切られる - リレーション毎にエッジはバケットに区切られる - パーティション内のリレーションは唯一 Data Model

- パーティションがP個とするとバケットは P^2個存在する - 1度にパーティションが重ならないバケットに対して学習が行 われる。あるパーティションにおける学習が終わり次第、別の パーティションとスワップされる - シングルマシン: ディスク - 分散学習: 各マシンにシャーディング Partition & Bucket

Partition & Bucket

Distributed Training

Distributed Training

0. すべてのマシンについて init_process_groupを実行し、以下を待つ - Parameter Server、Partition Serverを各マシンで起動し、 Rank 0 についてはLock Serverも起動 - Lock Server - バケットを管理する。各ワーカーから必要に応じてバ ケットを渡す Distributed Training

- Parameter Server - torch.distributed.{send/recv} を利用して実装 - Parameter Sharer - Parameter Serverに対して定期的にモデルにおける共 有パラメータ（nn.Moduleについてstate_dictで取得で きるもの）を同期する - https://github.com/facebookresearch/PyTorch-BigG raph/issues/34 - 他の通信を優先するため 100 updates/s またはデータサイズ 1GB/s の制約の中行われる Parameter Server

- Partition Server - Parameter Serverのクラスを用いて実装 - 各パーティションにおける Embeddingを管理。ワーカーから必 要に応じてパーティションをやり取りする （主にSwapping） - Partition Client - client = self._clients[part % len(self._clients)] Partition Server

- Rank 0のマシンをマスターとみなし、他のマシンはワーカーとして振 る舞う - マスター: 各ワーカーに対するソケットを作成し、ワーカーの 接続を待つ。ワーカーに対して各ワーカーの情報を送る - ワーカー: マスターへのソケットを作成し、マスターから他の ワーカーの情報を受け取る。すべてのワーカーと疎通確認を 行う - GILの開放も行う - https://github.com/pytorch/pytorch/blob/v1.0.1/torch/csr c/utils/auto_gil.h init_process_group

疎通後の各マシンのフロー : 1. Lock Serverに対してバケットの要求（ acquire_bucket）を行う - そのバケットが保有するパーティションはロックされる - 各マシンが互いに素なバケットで学習を行なう （Matrix FactorizationにおけるDSGDに似た話） 2. バケットからPartition Serverに対しパーティションをスワップ - 古いバケットにおけるパーティションのロックが開放される （release_bucket） Distributed Training

3. 学習 - 1 epoch: すべてのバケットが走査される - 各エッジは（分散）ファイルシステムから読み込まれ、各ロー カルスレッドに分けて学習される - Hogwild! - パラメータはワーカー間で非同期に更新される - OptimizerはAdaGradを利用 Distributed Training

1 epoch後: N. チェックポイントの作成 - 各ワーカーがPartition ServerからEmbeddingを取得しファイ ルに書き込む - torch.distributed.barrier Distributed Training

Single Machine Training

Single Machine Training - パーティションのスワップ（ディスクに対して）の効率化 - なるべくスワップする回数を減らしたい - Outside-in 1. あるインデックスN（0から）をもつ組み合わせをリストアップ - , <*,N> 2. 次に N+1について残りのリストを追加 3. 残りがなくなるまで続ける - 論文で紹介されているのはこれを反転させたもの （Inside-out）

Negative Sampling

- All negatives - 入力された正例に対して、同じパーティションに属するすべてのエンティティの中で可能 な限りすべてのエッジを負例とみなして学習を行う - 膨大な数の負例が生成されるので非効率 - Negative Sampling - パーティションを学習する際には入力がバッチに区切られる - 同じバッチにある他のエッジを負例として利用 - 他のバッチ（同じパーティションに属するエンティティ）から一様にサンプリング ⇒ これらを組み合わせて負例を生成 Negative Sampling

Evaluation

Evaluation Mean Reciprocal Rank (MRR): 正例の順位の逆数の平均 Hits@10: 負例との中で正例の Top10における割合

Evaluation 各パーティションサイズにおける Freebase datasetに対する10 epochの学習 左: シングルマシン, 右: 分散学習

About Me Software Engineer @Mercari ･ GitHub: @ynqa ･ Twitter: @yn_9a Go で Word Embeddings 書いてます ⇒ https://github.com/ynqa/wego

References

- PyTorch-BigGraph: A Large-scale Graph Embedding Framework - https://www.sysml.cc/doc/2019/71.pdf - Open-sourcing PyTorch-BigGraph for faster embeddings of extremely large graphs - https://ai.facebook.com/blog/open-sourcing-pytorch-biggraph-for-faster-embeddings-of-extremely- large-graphs - SysML 19: Adam Lerer, Pytorch-BigGraph: A Large Scale Graph Embedding System - https://www.youtube.com/watch?v=VPzwRQYM9-Y - PyTorch-BigGraph Navigation - https://torchbiggraph.readthedocs.io/en/latest/index.html References

References - DISTRIBUTED COMMUNICATION PACKAGE - https://pytorch.org/docs/stable/distributed.html - WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH - https://pytorch.org/tutorials/intermediate/dist_tuto.html - Knowledge Graph Embedding: A Survey of Approaches and Applications - https://persagen.com/files/misc/Wang2017Knowledge.pdf - A Comprehensive Survey of Graph Embedding: Problems, Techniques and Applications - https://arxiv.org/pdf/1709.07604.pdf - Holographic Embeddings of Knowledge Graphs - https://arxiv.org/pdf/1510.04935.pdf

- HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent - https://papers.nips.cc/paper/4390-hogwild-a-lock-free-approach-to-parallelizing-stochastic-gradien t-descent.pdf - Adaptive Subgradient Methods for Online Learning and Stochastic Optimization - http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf - Large-Scale Matrix Factorization with Distributed Stochastic Gradient Descent - http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.230.7682&rep=rep1&type=pdf References