Organizations: Institut de Recherche en Informatique de Toulouse, Toulouse, France. · Artificial and Natural Intelligence Toulouse Institute, Toulouse, France. · Centre National de la Recherche Scientifique, France. · Delft Institute of Applied Mathematics, Delft University of Technology, Delft, Netherlands. · Toulouse-INP (ENSEEIHT), Toulouse, France. · Department of Computer Science, University of Geneva, Geneva, Switzerland.
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.
Figures & tables
Figure 1: Illustration of MP aggregation for node v1 with feature vector h1 . The sketch is adapted from 47 .
Figure 2: An example of non-overlapping graph partitioning. The original graph (left) and a partition of the node set \pazocalV into P=3 disjoint subsets \pazocalV1,\pazocalV2,\pazocalV3 (right).
Figure 3: Sketch of the per-subdomain random node subsampling to create the coarse graph. The original graph (left) and the coarse graph \pazocalGC , formed by the selected nodes together with the edges inherited between them (right).
CIFAR10
AirfRANS
METR-LA
Learning task
Graph classification
Node regression
Spatiotemporal regression
GNN architecture
GCN 25
GraphSAGE 3
DCRNN 34
Number of layers
4
4
2
Layer width
146
64
32
Loss function
Cross-entropy
MSE
Masked MAE
Graph partitioner
Spectral 53
Morton ordering 36
Multilevel Graclus 10
Table 1: Summary of the benchmark problems and the associated experimental setup.
Figure 4: Example of a sample from CIFAR10 and its corresponding super-pixel graph 12 : the original image (left), the super-pixel nodes (middle left), the k -nearest-neighbor graph (middle right), and the resulting graph (right).
Figure 5: Example of flow fields around a NACA airfoil from the AirfRANS dataset: the pressure divided by specific mass (left) and the velocity field in the x -direction (right).
Figure 6: Example of a sample from the METR-LA dataset 34 : the sensor locations (left) and the corresponding graph partitioned into 5 subgraphs (right).
Figure 7: Convergence of AG2 and AG2m on the CIFAR10 (left) and AirfRANS (right) benchmark problems. The results are reported in terms of the validation accuracy (acc val ) for CIFAR10 and the validation loss ( \pazocalLval ) for AirfRANS.
Figure 8: Convergence of AG2m and DD-AG2m with varying KG for the CIFAR10. DD-AG2m is executed with P={2,3,5} (from left to right), and Kp=∣\pazocalD∣ .
Figure 9: Convergence of AG2m and DD-AG2m ( P=2 , KG=∣\pazocalD∣ ) with varying Kp for the CIFAR10 (left), for AirfRANS (middle) and for METR-LA (right).
Figure 10: Convergence of AG2m and DD-AG2m ( KG=Kp=∣\pazocalD∣ ) with increasing number of partitions P for the CIFAR10 (left), AirfRANS (middle), and METR-LA (right).
Figure 11: Convergence of AG2m and 2DD-AG2m ( KG=10,cf=2 ) with varying KC for the CIFAR10, P=2 (left), P=3 (middle), and P=5 (right).
Figure 12: Convergence of AG2m and 2DD-AG2m ( KG=10,KC=Kp=∣\pazocalD∣ ) on the AirfRANS (top) and METR-LA (bottom), with the coarsening factors cf=2 (left) and cf=4 (right) for different numbers of graph partitions P .
Figure 13: Convergence of AG2m , DD-AG2m ( KG=Kp=∣\pazocalD∣ , P=5 ) and 2DD-AG2m ( KG=10,KC=Kp=∣\pazocalD∣ , P=5 ) on the CIFAR10 (left, cf=2 ), AirfRANS (middle, cf=4 ), and METR-LA (right, cf=4 ).