Global Attention Improves Graph Networks Generalization

Abstract

This paper advocates incorporating a Low-Rank Global Attention (LRGA) module, a computation and memory efficient variant of the dot-product attention (Vaswani et al., 2017), to Graph Neural Networks (GNNs) for improving their generalization power. To theoretically quantify the generalization properties granted by adding the LRGA module to GNNs, we focus on a specific family of expressive GNNs and show that augmenting it with LRGA provides algorithmic alignment to a powerful graph isomorphism test, namely the 2-Folklore Weisfeiler-Lehman (2-FWL) algorithm. In more detail we: (i) consider the recent Random Graph Neural Network (RGNN) (Sato et al., 2020) framework and prove that it is universal in probability; (ii) show that RGNN augmented with LRGA aligns with 2-FWL update step via polynomial kernels; and (iii) bound the sample complexity of the kernel's feature map when learned with a randomly initialized two-layer MLP. From a practical point of view, augmenting existing GNN layers with LRGA produces state of the art results in current GNN benchmarks. Lastly, we observe that augmenting various GNN architectures with LRGA often closes the performance gap between different models.

References in corpus (18)

• Inductive Representation Learning on Large Graphs
• Neural Message Passing for Quantum Chemistry
• Fast Graph Representation Learning with PyTorch Geometric
• Learning Convolutional Neural Networks for Graphs
• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• A Survey on Graph Kernels
• Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks
• Residual Gated Graph ConvNets
• Benchmarking Graph Neural Networks
• Watch Your Step: Learning Node Embeddings via Graph Attention
• Graph Neural Tangent Kernel: Fusing Graph Neural Networks with Graph Kernels
• Geom-GCN: Geometric Graph Convolutional Networks
• Equivariance Through Parameter-Sharing
• Invariant and Equivariant Graph Networks
• Covariant Compositional Networks For Learning Graphs
• Attention Models in Graphs: A Survey
• Generalization and Representational Limits of Graph Neural Networks
• Relational Pooling for Graph Representations

Cited by in corpus (8)

• Improving Graph Neural Network Expressivity via Subgraph Isomorphism Counting
• Can Graph Neural Networks Count Substructures?
• Equivariant Subgraph Aggregation Networks
• VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization
• Frame Averaging for Invariant and Equivariant Network Design
• From Local Structures to Size Generalization in Graph Neural Networks
• Implicit SVD for Graph Representation Learning
• Fast Graph Learning with Unique Optimal Solutions