[Paper Review] Graph Embeddings at Scale

C. Bayan Bruss, Anish Khazane+4 more|arXiv (Cornell University)|Jul 3, 2019

Advanced Graph Neural Networks17 references4 citations

TL;DR

This paper proposes a distributed, parameter-server-based system for training graph embeddings at scale without graph partitioning. By using entity-based parameter servers and dynamic, size-constrained computational graphs, the system enables efficient GPU/CPU utilization and achieves high link prediction accuracy (90%) on a 50M-vertex heterogeneous graph and fast convergence on Friendster (68M vertices), with performance improving as worker nodes increase.

ABSTRACT

Graph embedding is a popular algorithmic approach for creating vector representations for individual vertices in networks. Training these algorithms at scale is important for creating embeddings that can be used for classification, ranking, recommendation and other common applications in industry. While industrial systems exist for training graph embeddings on large datasets, many of these distributed architectures are forced to partition copious amounts of data and model logic across many worker nodes. In this paper, we propose a distributed infrastructure that completely avoids graph partitioning, dynamically creates size constrained computational graphs across worker nodes, and uses highly efficient indexing operations for updating embeddings that allow the system to function at scale. We show that our system can scale an existing embeddings algorithm - skip-gram - to train on the open-source Friendster network (68 million vertices) and on an internal heterogeneous graph (50 million vertices). We measure the performance of our system on two key quantitative metrics: link-prediction accuracy and rate of convergence. We conclude this work by analyzing how a greater number of worker nodes actually improves our system's performance on the aforementioned metrics and discuss our next steps for rigorously evaluating the embedding vectors produced by our system.

Motivation & Objective

• To address the challenge of scaling graph embedding training to massive graphs with millions of vertices.
• To eliminate the need for costly graph partitioning in distributed training systems.
• To enable efficient, dynamic computation across worker nodes using entity-based parameter servers.
• To improve convergence speed and link prediction accuracy through asynchronous, distributed training.
• To evaluate the impact of increasing worker nodes on training performance and embedding quality.

Proposed method

• The system transforms the input graph into a fixed-size matrix of random walks, enabling row-wise partitioning across worker nodes without graph partitioning.
• Each worker node dynamically creates a TensorFlow computational graph per training iteration, limiting GPU memory usage.
• Entity-based parameter servers store and serve embeddings, allowing workers to fetch only required vertex embeddings via efficient indexing.
• Embedding updates are performed asynchronously using noise-contrastive estimation (NCE) loss, enabling scalable training on heterogeneous graphs.
• The architecture separates GPU-intensive training from CPU/RAM-based parameter storage, optimizing resource utilization.
• The system supports arbitrary numbers of worker nodes and scales efficiently by distributing gradient updates across all nodes.

Experimental results

Research questions

• RQ1 How can graph embedding training be scaled efficiently without graph partitioning?
• RQ2 What impact does increasing the number of worker nodes have on convergence speed and link prediction accuracy?
• RQ3 Can the system maintain high embedding quality while training on large-scale, heterogeneous graphs?
• RQ4 How does dynamic computational graph creation improve training efficiency and GPU utilization?
• RQ5 To what extent does asynchronous, entity-based parameter server communication improve scalability and performance?

Key findings

• The system achieved 90% link prediction accuracy on an internal heterogeneous graph with 50 million vertices and 280 million edges.
• On the Friendster dataset (68 million vertices), the system demonstrated fast convergence and high link prediction accuracy, with performance improving as worker count increased.
• Training with more worker nodes led to faster convergence and higher final accuracy, particularly after step 300, due to richer topographic signal integration.
• The system maintained high-quality embeddings even at scale, with no degradation in performance as the number of workers increased.
• The use of dynamic computational graphs and efficient indexing enabled effective GPU utilization while minimizing memory overhead.
• Asynchronous training with entity-based parameter servers allowed for globally regularized updates, improving embedding generalization.