Source-linked AI summary
PyTorch-BigGraph: A Large-scale Graph Embedding System
Adam Lerer, Ledell Wu, Jiajun Shen, Timothee Lacroix, Luca Wehrstedt, Abhijit Bose, Alex Peysakhovich
TL;DR
Large industrial graphs exceed the time and memory capacity of existing embedding systems. PBG uses adjacency-matrix partitioning and distributed execution to train embeddings at this scale, while supporting multi-relation graphs. It matches existing systems on common benchmarks and reports 88% lower memory consumption and 4× faster training for full Freebase under partitioning and eight-machine execution.
Problem
Industrial graphs with billions of nodes and trillions of edges create severe embedding time and memory challenges; two billion nodes with 100 float parameters per node require 800GB.
Method
PBG trains one adjacency-matrix partition at a time, swapping embeddings to disk or distributing execution across machines, while supporting multi-entity, multi-relation graphs.
Results
PBG matches existing systems on common benchmarks; on full Freebase, partitioning reduces memory consumption by 88% and eight-machine execution speeds training by 4×.
Takeaways & Limitations
PBG enables graph embeddings for graphs with billions of nodes and trillions of edges on a single machine or in distributed environments.
Takeaways & Limitations
Embedding quality may be more sensitive to parallelization when relation counts or degree skew are high, or when using operators such as ComplEx.
Abstract
from arXiv · showhide
Graph embedding methods produce unsupervised node features from graphs that can then be used for a variety of machine learning tasks. Modern graphs, particularly in industrial applications, contain billions of nodes and trillions of edges, which exceeds the capability of existing embedding systems. We present PyTorch-BigGraph (PBG), an embedding system that incorporates several modifications to traditional multi-relation embedding systems that allow it to scale to graphs with billions of nodes and trillions of edges. PBG uses graph partitioning to train arbitrarily large embeddings on either a single machine or in a distributed environment. We demonstrate comparable performance with existing embedding systems on common benchmarks, while allowing for scaling to arbitrarily large graphs and parallelization on multiple machines. We train and evaluate embeddings on several large social network graphs as well as the full Freebase dataset, which contains over 100 million nodes and 2 billion edges.
1 INTRODUCTION
Graph embeddings provide useful node features but face severe time and memory barriers on industrial-scale graphs. PBG addresses these barriers through partitioned, distributed training while supporting multiple entity and relation types.
- Motivation: Industrial graphs can contain billions of nodes and trillions of edges, making embedding speed and memory major challenges.A two-billion-node model with 100 float parameters per node would require 800GB just for parameters.
- Approach: PBG uses adjacency-matrix block decomposition, training one bucket at a time and swapping or distributing partition embeddings.This design reduces memory usage and supports execution on one machine or across multiple machines.
- Approach: PBG improves training efficiency through negative sampling that combines uniform and data-based sampling and reuses negatives within batches.The reuse strategy reduces memory bandwidth requirements.
- Approach: PBG supports multi-entity, multi-relation graphs with per-relation edge weights and relation operators.Its distributed model also uses parameter-server infrastructure for global parameters and feature embeddings.
- Results: PBG matches existing embedding systems on Freebase, LiveJournal, and YouTube benchmarks.The system was also evaluated on larger graphs, including full Freebase and a large Twitter graph.
- Results: 88% lower memory consumption and 4× faster training were achieved for full Freebase with partitioning and eight-machine distributed execution.The embedded Freebase graph contains 121 million entities and 2.4 billion edges; partitioning caused only a small quality degradation.
2 RELATED WORK
Prior work spans multi-relation embeddings, distributed and specialized scalable algorithms, distributional word representations, and graph convolutional networks. PBG relates most directly to efforts that scale embedding or matrix-factorization models to large data.
- Multi-relation embeddings: Multi-relation embedding models represent entities with base vectors and relation-specific transformations, often targeting knowledge-base graphs.Related work also models some entities as bags of other entities rather than assigning them explicit embeddings.
- Scalable algorithms: Scalable graph-learning research includes multilevel methods, distributed embedding systems, and specialized large-graph algorithms for SVD and k-means.Large embedding systems have also produced gains in e-commerce and other applications.
- Distributional semantics: Word2vec enabled word embeddings to scale to larger corpora, with later distributed systems showing economic value from ingesting still larger datasets.This work belongs to the broader literature on distributional semantics in natural language processing.
- Related scalable methods: Parallel machine-learning algorithms and scalable matrix-factorization methods provide closely related foundations for distributed embedding systems.Matrix factorization has also been widely successful in recommender systems.
- Graph convolutional networks: Graph convolutional networks construct embeddings for large-scale applications, but typically operate on graphs whose nodes already have features.The GCN problem differs from the one addressed by PBG; combining the approaches is identified as future work.
3 MULTI-RELATION EMBEDDINGS
PBG represents multi-relation graphs with entity and relation parameters, scores transformed endpoint vectors, and supports several relation operators and training objectives. It constructs negatives by corrupting edges, using sampling choices tailored to sparse and multi-entity graphs.
- A multi-relation graph contains entities, relations, and directed edges, with relations optionally constrained to source and destination entity types.
- The model represents each entity and relation type with parameter vectors and updates them using minibatch stochastic gradient descent with Adagrad.
- PBG scores transformed source and destination vectors using relation-specific operators and similarity functions.Supported operators include linear transformation, translation, and complex multiplication; supported similarities include dot product and cosine.
- PBG can train RESCAL, DistMult, TransE, and ComplEx models through combinations of scoring functions and relation operators.
- Negative sampling mixes nodes drawn from training-data prevalence with uniformly sampled nodes to avoid undesirable behavior from either distribution alone.In multi-entity graphs, negatives are restricted to the entity type allowed by the relation.
- PBG optimizes a margin-based ranking objective by comparing each positive edge with negatives formed by replacing either its source or destination.Logistic and softmax losses can also be used to reproduce certain graph embedding models.
4 TRAINING AT SCALE
PBG scales training by partitioning entities and edges into memory-sized buckets, then swapping or distributing partitions for execution. It further reduces communication and memory costs through parallel bucket scheduling and batched negative sampling.
- PBG partitions models too large for one machine, enabling either disk-based single-machine training or distributed execution.
- Each partitioned entity type is split into P parts, and edges are grouped into buckets by source and destination partitions.This yields P^2 buckets when both endpoint types are partitioned and P buckets when only one is partitioned.
- Training iterates through buckets while loading the required partitions, with embeddings from each partition swapped from disk and edges subdivided across threads.
- Partitioning changes negative comparisons and groups edges by partition, preserving convergence guarantees but potentially slowing convergence.Negatives are drawn from the same source or destination partition as the positive edge.
- An inside-out bucket order improves single-machine embedding performance while minimizing disk swaps, and ensures previously trained partitions are available.
- Distributed training locks embedding partitions, trains buckets with disjoint partitions in parallel, and uses asynchronous parameter-server updates for shared parameters.The scheme can train in parallel on up to P/2 machines.
- Batched negative sampling reuses sampled embeddings across edge chunks, reducing memory fetches while computing many negative scores with matrix multiplication.For each batch of B positive edges, only 3B embeddings are fetched and 3BBn dot-product scores are computed.
- Grouping edges by relation accelerates the linear relation operator by expressing its transformation as a matrix multiplication.
5 EXPERIMENTS
PBG is evaluated on social-network and knowledge-graph tasks, where it matches competing embedding systems while supporting large-scale partitioned and distributed training. Experiments on full Freebase show near-preserved accuracy alongside reduced memory usage and faster multi-machine training.
- Experimental scope: PBG is evaluated on link prediction and node-attribute prediction across social networks and the Freebase knowledge graph.The social-network experiments include LiveJournal, Twitter, and YouTube; Freebase evaluates knowledge-graph link prediction.
- Partitioning and scaling: On one machine, peak memory decreases almost linearly with partition count, while training time rises because of additional I/O.With multiple machines, training time decreases as machines increase, but I/O and incomplete occupancy add overhead.
- Full Freebase: Freebase link-prediction accuracy remains nearly identical through four machines, while eight machines reduce MRR from 0.171 to 0.163.The reported degradation occurs under the highest parallelization condition.
- Partitioning and scaling: Twitter shows more linear training-time scaling as the graph is partitioned and trained in parallel.The comparison is made against the Freebase learning curves.
6 CONCLUSION
PBG scales multi-entity, multi-relation graph embeddings through bucketed adjacency decomposition, supporting reduced memory use and parallelization. Its embeddings match existing systems, while quality can be more sensitive to parallelization for complicated models; the largest benefits are expected on graphs 1–2 orders of magnitude larger than evaluated datasets.
- PBG supports multi-entity, multi-relation graphs with per-relation configuration such as edge weight and relation operator choice.
- PBG decomposes the adjacency matrix into N buckets and trains one bucket at a time to save memory and enable parallelization.
- 88% lower memory consumption on Freebase was achieved without degrading embedding quality, while 8-machine distributed execution made training 4× faster.
- Embedding quality may be more sensitive to parallelization when relation counts or degree skew are high, or when operators such as ComplEx are used.
- The architecture’s largest benefits are expected on graphs 1–2 orders of magnitude larger than the publicly available datasets evaluated.