Learning Large Graph Property Prediction via Graph Segment Training

TL;DR

This paper introduces Graph Segment Training (GST), a memory-efficient divide-and-conquer framework for large graph property prediction, incorporating techniques to handle embedding staleness and achieve competitive accuracy.

Contribution

The paper proposes GST, a novel framework that enables large graph learning with constant memory, and introduces methods to address embedding staleness and input distribution shifts.

Findings

GST-EFD is memory-efficient and fast.

GST-EFD slightly outperforms full graph training in accuracy.

Effective handling of embedding staleness improves performance.

Abstract

Learning to predict properties of large graphs is challenging because each prediction requires the knowledge of an entire graph, while the amount of memory available during training is bounded. Here we propose Graph Segment Training (GST), a general framework that utilizes a divide-and-conquer approach to allow learning large graph property prediction with a constant memory footprint. GST first divides a large graph into segments and then backpropagates through only a few segments sampled per training iteration. We refine the GST paradigm by introducing a historical embedding table to efficiently obtain embeddings for segments not sampled for backpropagation. To mitigate the staleness of historical embeddings, we design two novel techniques. First, we finetune the prediction head to fix the input distribution shift. Second, we introduce Stale Embedding Dropout to drop some stale embeddings during training to reduce bias. We evaluate our complete method GST-EFD (with all the techniques together) on two large graph property prediction benchmarks: MalNet and TpuGraphs. Our experiments show that GST-EFD is both memory-efficient and fast, while offering a slight boost on test accuracy over a typical full graph training regime.