Source-linked AI summary
Two-level domain-decomposition AdaGrad method for scalable training of graph neural networks
Laurynas Varnas, Julien Herrmann, Alexander Heinlein, Serge Gratton, Alena Kopaničáková
TL;DR
Efficient distributed GNN training is limited by message-passing costs, memory requirements, and communication overhead. The paper introduces DD-AG2m, which combines global and partition-level curvature-informed AdaGrad updates, and 2DD-AG2m, which uses a coarse graph for cheaper global steps. Across graph classification, node-level regression, and spatiotemporal forecasting, the methods require fewer optimization steps for equivalent predictive performance and improve performance at fixed computational cost.
Problem
Distributed GNN training is challenging because message passing couples graph nodes, producing expensive optimization, high memory requirements, and substantial communication overhead.
Method
DD-AG2m alternates AG2m optimization on the global and partitioned graphs, while 2DD-AG2m performs global steps on a randomly subsampled coarse graph.
Results
Across graph classification, node-level regression, and spatiotemporal forecasting, the DD methods require up to 4–8× fewer optimization steps than AG2m for the same predictive performance.
Takeaways & Limitations
The proposed DD methods improve GNN predictive performance for a fixed computational cost while exposing additional parallelism through graph partitioning.
Takeaways & Limitations
A distributed-memory implementation is still needed to translate the observed computational-cost reduction into a wall-clock speedup.
Abstract
from arXiv · showhide
Graph neural networks (GNNs) have emerged as a powerful framework for learning from graph-structured data. However, their efficient training remains challenging, particularly in distributed computing environments. This challenge arises from the use of message passing, which couples all graph nodes, leading to expensive optimization steps, high memory requirements, and substantial communication overhead. To alleviate these limitations, we propose a novel domain-decomposition (DD) variant of AG2m, an AdaGrad method enhanced with second-order curvature information and momentum, denoted by DD-AG2m. The proposed DD-AG2m alternates between AG2m optimization on the original (global) graph and AG2m optimization on the partitioned graphs. To incorporate global information at reduced cost, we further introduce a two-level variant (2DD-AG2m) that performs global optimization steps on a coarse graph obtained by randomly subsampling nodes within each subdomain. Numerical experiments spanning graph classification, node-level regression, and spatiotemporal forecasting tasks demonstrate that the proposed DD methods reduce the computational cost required to achieve the same predictive performance by a factor of 4-8. Moreover, for the fixed computational cost, they improve the predictive performance of GNNs by up to 22% compared with the baseline AG2m.
1 Introduction
GNN training at scale is difficult because message passing creates irregular computation, memory demands, and distributed communication costs. The paper proposes domain-decomposition AdaGrad methods that use local graph optimization, with a coarse-graph two-level variant to retain global information more cheaply.
- Motivation: Distributed GNN training is hindered by irregular workloads, memory access, synchronization, and feature-gradient communication costs.Message-passing dependencies can increase synchronization costs and dominate training time.
- Existing approaches: Reduced computation graphs, sampling, clustering, and partitioning limit neighborhood computations and enable potentially parallel subgraph processing.Existing approaches include node- and layer-wise sampling, Cluster-GCN, and GraphSAINT.
- Existing approaches: Local partition training can retain residual error when cross-partition dependencies are ignored, motivating methods that coordinate local updates globally.LLCG uses periodic averaging and global server corrections, while periodic averaging alone may leave residual error.
- Domain decomposition: Domain-decomposition methods solve smaller subproblems concurrently while coordinating information exchange to recover a globally coherent solution.This design provides a natural framework for parallelization and scalable computation.
- Proposed methods: DD-AG2m partitions the input graph, performs independent local optimization with temporary shared-model copies, and aggregates local corrections.The aggregation follows the spirit of additive Schwarz methods.
- Proposed methods: 2DD-AG2m performs global optimization steps on a coarse graph formed by randomly subsampling nodes within each subdomain.The approach treats partitioning as a nonlinear preconditioning mechanism for AdaGrad while exposing additional parallelism.
2 Graph Neural Networks (GNN)
GNNs map attributed graphs to graph-, node-, or node-time predictions through message passing and task-specific readout functions. Training minimizes empirical risk, but full-graph message passing creates irregular workloads, growing memory requirements, and large-graph memory constraints.
- GNN model: A GNN maps an attributed graph to graph-level, node-level, or node-time predictions through parameterized message-passing and readout functions.The paper evaluates graph classification, node regression, and spatiotemporal forecasting settings.
- Message passing: Each node updates its representation by aggregating information from its local neighborhood, with graph connectivity defining computational dependencies.After L layers, a node representation depends only on its L-hop neighborhood.
- Message passing: Message-passing parameters are shared across nodes, and permutation-invariant aggregation makes node representations equivariant to node-order permutations.The framework permits neighborhoods of unrestricted size.
- GNN training: Training minimizes empirical risk over input graphs and supervised targets using a task-specific loss comparing predictions with targets.The optimization problem is expressed as minimizing L(θ) over model parameters.
- Training challenges: Non-uniform node degrees induce irregular workloads and memory-access patterns that can reduce accelerator utilization and complicate parallel execution.The cost depends on sample count, graph size, and graph topology.
- Training challenges: O(L |V| d′) memory growth with depth reflects the intermediate representations and message-passing dependencies needed for gradient-based optimization.Backpropagation requires storing or recomputing intermediate node representations across layers.
- Training challenges: Full-graph processing may exceed device memory for very large graphs, while subgraph restriction changes the loss function and its gradient.The proposed DD methods use subgraphs to construct local corrections and coordinate them through global optimization steps.
3 Two-level DD-based AdaGrad for GNNs
The method combines curvature-informed AdaGrad with momentum and domain decomposition to train GNNs through global and partitioned-graph optimization. A two-level variant replaces some global steps with coarse-graph optimization to transfer global information at lower cost.
- Single-level curvature-informed AdaGrad with momentum: AG2m combines AdaGrad with curvature information and momentum as the computational engine for DD-AG2m and 2DD-AG2m.It uses adaptive coordinate-wise bounds, curvature-based scaling, and bounded momentum updates.
- Global and subdomain minimization problems: DD-AG2m partitions each graph into disjoint subgraphs, defines subdomain losses, and maintains subdomain parameters with the same structure as the global model.Each induced subgraph removes inter-partition edges, while model parameters remain structurally unchanged across subdomains.
- DD-AG2m algorithm: Each DD-AG2m outer iteration first performs KG AG2m steps on the full graph, then applies Kp AG2m steps independently to each subdomain.Subdomain optimization starts from the post-global parameters and uses inherited AdaGrad weights.
- DD-AG2m algorithm: Subdomain corrections are averaged and applied through a single global reduction to update the shared GNN parameters.The algorithm then proceeds to the next outer iteration.
- DD-AG2m algorithm: DD-AG2m exposes parallel subgraph computation, but its inherently serial KG global steps remain the main computational bottleneck.The method accelerates AG2m convergence through subgraph computations while retaining global optimization stages.
- Two-level DD-based AdaGrad: The two-level method replaces several full-graph optimization steps with coarse-graph steps that transfer global information at reduced cost.The coarse graph retains randomly subsampled nodes selected independently within each subdomain, with at least one retained node per subdomain; its edges are inherited from the original graph.
4 Numerical examples and implementation details
The evaluation covers graph classification, physics-based node-level regression, and spatiotemporal traffic forecasting, using diverse datasets, GNN architectures, partitioning strategies, and implementation settings.
- Benchmark problems: Three benchmark problems assess the methods across graph classification, physics-based regression, and spatiotemporal traffic forecasting.The benchmarks are CIFAR10, AirfRANS, and METR-LA.
- CIFAR10 graph classification: CIFAR10 uses SLIC-derived super-pixel graphs and a four-layer GCN trained with cross-entropy loss.Graphs are partitioned with spectral partitioning.
- AirfRANS regression: AirfRANS contains steady turbulent flow simulations around NACA airfoils, evaluated in a scarce-data setting with 200 training and 200 testing simulations.Each simulation is uniformly subsampled to 32,000 nodes before sparse geometric graph construction.
- METR-LA forecasting: METR-LA formulates traffic prediction as node-level spatiotemporal regression on a fixed 207-sensor road-network graph.Inputs contain 12 historical time steps and predictions cover the following 12 steps; the model uses masked MAE.
- Partitioning and implementation: Graph partitions remain fixed during training, whereas the coarse graph is regenerated each epoch through per-subdomain random node subsampling.The coarsening factor specifies the fraction of nodes retained in each subdomain.
5 Numerical results
The experiments evaluate convergence, computational cost, partition scalability, and coarse-level choices for AG2m, DD-AG2m, and 2DD-AG2m across three benchmark problems. DD-based methods generally achieve more accurate models at a given computational cost, while coarse steps can improve early convergence and reduce expensive global work.
- Computational cost: A subdomain step costs 1/P and a coarse-graph step costs 1/cf, while global full-graph steps have unit cost.The accumulated cost uses KG, Kp, and KC, the numbers of global, subdomain, and coarse-level steps per outer iteration.
- AG2m convergence: Momentum improves initial convergence, while AG2m reaches comparable final validation accuracy to AG2 on CIFAR10.AG2m also attains lower validation loss than AG2 for comparable computational cost.
- DD-AG2m convergence: 1, 176, and 352 global steps produce different trade-offs: KG = 1 accelerates initial training, whereas KG ∈ {176, 352} yields higher final validation accuracy.This sensitivity study uses CIFAR10, Kp = 352, and P ∈ {2, 3, 5}.
- DD-AG2m convergence: Kp = 2|D| provides the best or comparable performance across all datasets, while substantially larger values may overfit subgraph models.The comparison uses P = 2, KG = |D|, and Kp ∈ {100, |D|, 2|D|}.
- Partition scalability: DD-AG2m produces more accurate models than AG2m at a given computational cost, without convergence deterioration as P increases.The comparison covers CIFAR10, AirfRANS, and METR-LA with increasing numbers of partitions.
- 2DD-AG2m convergence: KC = |D| and KC = 2|D| yield comparable 2DD-AG2m performance, indicating that coarse steps can be reduced without significantly affecting the final solution.The coarse phase replaces expensive full-graph steps with cheaper coarsened-graph steps.
- 2DD-AG2m convergence: Increasing cf from 2 to 4 generally improves 2DD-AG2m convergence on AirfRANS and METR-LA, especially with more partitions.The largest performance and accuracy improvements occur on AirfRANS when larger coarsening factors are combined with more partitions.
- Method comparison: At P = 5, both DD methods outperform AG2m at fixed computational cost, and 2DD-AG2m converges faster early than DD-AG2m.On CIFAR10, 2DD-AG2m also reaches notably higher final validation accuracy; coarse steps exchange global information at a fraction of full-graph cost.
6 Conclusion
The paper proposes DD-AG2m and 2DD-AG2m for scalable GNN training and evaluates them across graph classification, node regression, and spatiotemporal forecasting. They reduce optimization steps for equal predictive performance, improve performance at equal steps, and remain stable as partition count increases.
- Conclusion: DD-AG2m and 2DD-AG2m combine partition-level optimization with full-graph or randomly subsampled coarse-graph optimization.Both methods are built on curvature-informed AdaGrad with momentum and are evaluated with different GNN architectures and partitioning strategies.
- Conclusion: 4–8× fewer optimization steps are required to achieve the same predictive performance as AG2m.This result spans graph classification, node-level regression, and spatiotemporal forecasting tasks.
- Conclusion: Up to 22% higher predictive performance is achieved at the same number of optimization steps as AG2m.Performance remained stable as the number of graph partitions increased.
- Future work: A distributed-memory implementation, overlapping subdomains, alternative coarsening strategies, multilevel extensions, and asynchronous variants remain future work.The authors specifically plan to translate computational-cost reductions into wall-clock speedups through distributed-memory implementation.