Source-linked AI summary
Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks
Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, Cho-Jui Hsieh
TL;DR
Large-scale GCN training is constrained by layer-dependent computational costs and memory demands from node embeddings. Cluster-GCN uses graph-clustered subgraphs for SGD batches, achieving efficient training and strong accuracy on large and deep GCN tasks. Its main limitation is that partitioning removes some links and can bias gradient estimates.
Problem
Large-scale GCN training remains challenging because existing SGD methods can have exponentially increasing computational complexity with depth or large memory requirements.
Method
Cluster-GCN forms SGD batches from dense graph-clustered subgraphs, restricts neighborhood search within each subgraph, and uses stochastic multiple partitions.
Results
Cluster-GCN trains very deep GCNs efficiently on large graphs and achieves state-of-the-art test F1 scores on PPI and Reddit.
Takeaways & Limitations
Cluster-GCN enables training deeper GCNs on large-scale graphs with low memory use and competitive or improved prediction accuracy.
Takeaways & Limitations
Partitioning removes some graph links and can produce biased full-gradient estimates because clusters may differ from the original data distribution.
Abstract
from arXiv · showhide
Graph convolutional network (GCN) has been successfully applied to many graph-based applications; however, training a large-scale GCN remains challenging. Current SGD-based algorithms suffer from either a high computational cost that exponentially grows with number of GCN layers, or a large space requirement for keeping the entire graph and the embedding of each node in memory. In this paper, we propose Cluster-GCN, a novel GCN algorithm that is suitable for SGD-based training by exploiting the graph clustering structure. Cluster-GCN works as the following: at each step, it samples a block of nodes that associate with a dense subgraph identified by a graph clustering algorithm, and restricts the neighborhood search within this subgraph. This simple but effective strategy leads to significantly improved memory and computational efficiency while being able to achieve comparable test accuracy with previous algorithms. To test the scalability of our algorithm, we create a new Amazon2M data with 2 million nodes and 61 million edges which is more than 5 times larger than the previous largest publicly available dataset (Reddit). For training a 3-layer GCN on this data, Cluster-GCN is faster than the previous state-of-the-art VR-GCN (1523 seconds vs 1961 seconds) and using much less memory (2.2GB vs 11.2GB). Furthermore, for training 4 layer GCN on this data, our algorithm can finish in around 36 minutes while all the existing GCN training algorithms fail to train due to the out-of-memory issue. Furthermore, Cluster-GCN allows us to train much deeper GCN without much time and memory overhead, which leads to improved prediction accuracy---using a 5-layer Cluster-GCN, we achieve state-of-the-art test F1 score 99.36 on the PPI dataset, while the previous best result was 98.71 by [16]. Our codes are publicly available at https://github.com/google-research/google-research/tree/master/cluster_gcn.
1 INTRODUCTION
GCN training is difficult at scale because node dependence makes computation and memory grow substantially, motivating scalable algorithms that balance memory, epoch time, and convergence. Cluster-GCN uses graph clustering to form efficient mini-batches and reports strong scalability and accuracy.
- GCN training is slow and memory-intensive because each node’s loss depends on many other nodes and back-propagation stores embeddings across the computation graph.
- Existing algorithms trade off memory, time per epoch, and convergence speed, with full-batch methods requiring O(NFL) memory and mini-batch methods incurring neighborhood-expansion overhead.
- Cluster-GCN exploits graph clustering to create batches with more within-batch links, requiring O(bFL) memory for embeddings in the current batch.
- Cluster-GCN combines clustering-based batching with stochastic multi-clustering to improve convergence while reducing memory and computational costs.
- Amazon2M contains more than 2 million nodes and 61 million edges, providing a larger graph for evaluating GCN scalability.
2 BACKGROUND
GCN layers propagate node representations by repeatedly mixing neighboring embeddings, transforming them, and applying nonlinear activations. The model operates on a sparse graph with node features and learned layer-specific weight matrices.
- An L-layer GCN constructs each node embedding by mixing embeddings from neighboring nodes in the previous layer.
- The graph is represented by an N × N sparse adjacency matrix A, while X ∈ R^N×F stores F-dimensional features for all N nodes.
- Each layer applies the normalized, regularized adjacency matrix and a learned feature-transformation matrix, followed by an activation such as element-wise ReLU.
- For semi-supervised node classification, the weight matrices are learned by minimizing a loss over labeled nodes, commonly using cross-entropy.
- Table 1 analyzes training complexity using L, N, ∥A∥0, F, batch size b, and sampled-neighbor count r, while separating parameter and embedding-storage memory.
3 PROPOSED ALGORITHM
Cluster-GCN makes mini-batch GCN training efficient by clustering nodes into dense subgraphs, restricting neighborhood computation, and addressing partition-induced bias with stochastic multiple partitions. It achieves linear-in-depth embedding computation, batch-sized embedding memory, and deeper-GCN training through diagonal enhancement.
- Motivation: Mini-batch SGD suffers from exponential neighborhood-expansion cost as GCN depth increases, while full-batch training requires O(NFL) embedding memory.Random sampling computes O(bdL) embeddings per batch and O(NdLF2) time per epoch under sparse-neighborhood assumptions.
- Vanilla Cluster-GCN: Cluster-GCN maximizes embedding utilization by constructing batches with many within-batch edges using graph clustering.Embedding utilization is equivalent to the number of within-batch edges, motivating partitions with more within-cluster than between-cluster links.
- Vanilla Cluster-GCN: Restricting neighborhood aggregation to a clustered batch reduces neighborhood expansion while preserving repeated reuse of the same subgraph across GCN layers.The method focuses computation on within-cluster neighbors rather than searching the full graph.
- Complexity: Cluster-GCN requires O(bLF) embedding memory and O(∥A∥0F + NF2) time per epoch, with O(bL) embeddings computed per batch.It loads only a subgraph and the batch embeddings instead of the full graph and all node-layer embeddings.
- Stochastic Multiple Partitions: Stochastic multiple clustering samples several clusters and restores their between-cluster links, reducing partition bias and variance across batches.Different cluster combinations are selected across epochs, and experiments report improved convergence when multiple clusters form a batch.
- Deep GCNs: Diagonal enhancement gives greater weight to previous-layer representations, helping deep GCNs maintain neighborhood information and achieve stronger accuracy.The modified normalization strategy was reported to support deep GCN training and state-of-the-art performance.
4 EXPERIMENTS
Experiments evaluate Cluster-GCN across public graph datasets and show strong speed, memory, scalability, convergence, and deep-GCN performance.
- Experimental setup: Cluster-GCN is evaluated on multi-label and multi-class classification across four public datasets, including Reddit and the larger Amazon2M graph.Amazon2M contains over 2 million nodes and 61 million edges.
- Training performance: Cluster-GCN is fastest on PPI and Reddit across GCNs with 2, 3, and 4 layers.The comparison plots training time against validation F1 accuracy.
- Training performance: PyTorch sparse tensor performance can make Cluster-GCN slower than VRGCN on some Amazon settings despite being faster for the 3-layer case.Amazon uses an identity feature matrix, making sparse matrix operations dominant; TensorFlow is faster than PyTorch in the reported benchmark.
- Memory usage: 308MB versus 2064MB: Cluster-GCN uses substantially less memory than VRGCN for a 4-layer Reddit GCN with hidden dimension 512.The paper attributes VRGCN’s higher memory use to storing historical embeddings, while GraphSAGE faces exponential neighborhood growth.
- Amazon2M: Over 2 million nodes and 61 million edges: Amazon2M tests scalability beyond the approximately 200K-node Reddit dataset.The graph is constructed from Amazon co-purchasing networks, with products as nodes.
- Training deeper GCNs: Exponential versus linear growth: VRGCN’s running time increases exponentially with depth, whereas Cluster-GCN’s increases linearly on PPI.Both methods are compared over 200 epochs.
- Training deeper GCNs: 2 to 5 layers: accuracy increases for all methods on PPI, while the first three methods fail to converge within 200 epochs at 7 or 8 layers.With diagonal enhancement technique (11), convergence improves significantly and similar accuracy is achieved for an 8-layer GCN.
5 CONCLUSION
ClusterGCN is presented as a fast, memory-efficient training algorithm that enables very deep GCNs on large graphs and achieves state-of-the-art test F1 scores.
- Conclusion: ClusterGCN trains very deep GCNs on large-scale graphs while using limited time and memory.On a graph with over 2 million nodes, training takes less than an hour and uses around 2G memory.
- Conclusion: State-of-the-art test F1 scores are achieved on the PPI and Reddit datasets using deeper GCNs.The conclusion reports these results as enabled by the proposed approach.
6 MORE DETAILS ABOUT THE EXPERIMENTS
The experiments use specific dataset, preprocessing, normalization, and memory-measurement settings to support reproducibility. Cluster-GCN’s graph-clustering preprocessing adds a small, reusable cost before training.
- Experimental settings: The experiments use PPI, Reddit, and Amazon datasets, with Amazon evaluated in an inductive setting where training excludes testing nodes.Software versions, METIS, and dataset sources are specified for reproducibility.
- Implementation details: Precomputing AX uses each node’s exact 1-hop neighborhood and avoids expensive neighborhood searches in the first GCN layer.This implementation strategy is adopted from previous work.
- Implementation details: When multiple clusters are selected, between-cluster links are restored and the combined adjacency matrix is re-normalized to maintain embedding numerical ranges.Experiments find renormalization helpful.
- Implementation details: Inductive experiments partition only the training-node adjacency matrix, then re-normalize the partitioned matrix while using an all-node matrix for evaluation.Feature normalization is also conducted, and memory is measured with framework-specific GPU or TensorFlow counters.
- Preprocessing: Reported training times exclude preprocessing, although methods may incur algorithm-specific preprocessing such as graph clustering.The paper separately reports graph-clustering and preprocessing times in Table 13.
- Preprocessing: METIS graph clustering takes a small portion of preprocessing time and is performed once to create node partitions reusable across later training processes.The paper characterizes this as a small extra cost on large datasets.