PaperlayerView exact revision on arXiv

Source-linked AI summary

Global Self-Attention as a Replacement for Graph Convolution

Md Shamim Hussain, Mohammed J. Zaki, Dharmashankar Subramanian

arXiv:2108.03348v3cs.LG

TL;DR

Graph learning needs a way to handle arbitrary structural information without relying on fixed local aggregation. EGT adds residual edge channels to a transformer and uses global self-attention; it outperforms convolutional GNNs across experiments and reaches a new state of the art on PCQM4Mv2, with an acknowledged theoretical generalization caveat for absolute positional coordinates.

  • Problem

    Graph structure is arbitrary, positional encodings alone may not represent node positions adequately, and existing transformer adaptations do not directly process edge information for general graph prediction.

  • Method

    EGT extends the transformer with residual edge channels that hierarchically process structural information while using global self-attention for dynamic, unconstrained aggregation.

  • Results

    EGT achieves superior performance over convolutional GNNs across medium-scale, large-scale, and transfer-learning settings and sets a new state of the art on PCQM4Mv2.

  • Takeaways & Limitations

    The experiments indicate that fixed local convolutional aggregation is not an essential inductive bias for graph neural networks.

  • Takeaways & Limitations

    Absolute SVD-based coordinates may theoretically hamper generalization because they depend on an input graph’s reference frame, although no detrimental effect was observed in practice.

Abstract

from arXiv · show

We propose an extension to the transformer neural network architecture for general-purpose graph learning by adding a dedicated pathway for pairwise structural information, called edge channels. The resultant framework - which we call Edge-augmented Graph Transformer (EGT) - can directly accept, process and output structural information of arbitrary form, which is important for effective learning on graph-structured data. Our model exclusively uses global self-attention as an aggregation mechanism rather than static localized convolutional aggregation. This allows for unconstrained long-range dynamic interactions between nodes. Moreover, the edge channels allow the structural information to evolve from layer to layer, and prediction tasks on edges/links can be performed directly from the output embeddings of these channels. We verify the performance of EGT in a wide range of graph-learning experiments on benchmark datasets, in which it outperforms Convolutional/Message-Passing Graph Neural Networks. EGT sets a new state-of-the-art for the quantum-chemical regression task on the OGB-LSC PCQM4Mv2 dataset containing 3.8 million molecular graphs. Our findings indicate that global self-attention based aggregation can serve as a flexible, adaptive and effective replacement of graph convolution for general-purpose graph learning. Therefore, convolutional local neighborhood aggregation is not an essential inductive bias.

1 INTRODUCTION

Graph learning commonly relies on localized convolutional aggregation, but graph structure makes positional encoding and edge-feature integration difficult. EGT addresses these challenges with residual edge channels and global, dynamic self-attention that can aggregate distant nodes directly.

  • Graph-structured data are difficult to process because arbitrary node interconnectivity increases learning complexity.
  • Convolutional GNNs update each node from its neighbors, imposing a localized aggregation pattern at every layer.
  • Global self-attention forms input-dependent aggregation weights, allowing nodes to aggregate information globally rather than through a fixed local window.
  • Residual edge channels let EGT directly process arbitrary structural information, including edge features and directed or weighted graph variants.
  • Global self-attention can learn to aggregate distant graph regions in one step, with a pattern that adapts to each input graph.

2 RELATED WORK

Prior graph attention methods often restrict attention to local neighborhoods, while other approaches use attention for heterogeneous-edge composition. EGT instead uses attention as global aggregation for homogeneous graphs.

  • GAT and GT constrain self-attention to each node’s local neighborhood, resembling graph convolution and local message passing.
  • GTN and related methods use attention to combine heterogeneous edges across multiple hops, whereas EGT focuses on global aggregation in homogeneous graphs.

3 NETWORK ARCHITECTURE

EGT extends the transformer with edge channels that represent, update, and output pairwise structural information alongside node embeddings. Its attention uses edge information to form dynamic, globally connected aggregation patterns, with centrality scaling and optional positional encodings for graph learning.

  • 3.1 Preliminaries: Transformer layers combine multihead self-attention with pointwise feed-forward transformations, using layer normalization and residual connections in the EGT architecture.The original transformer’s node channels process node embeddings, while EGT introduces parallel edge channels for pairwise structure.
  • 3.2 Edge-augmented Graph Transformer (EGT): EGT adds N×N edge embeddings for node pairs, formed from graph structural matrices and edge features and updated at every layer.The model outputs edge embeddings that support structural predictions such as edge labeling and link prediction.
  • 3.2 Edge-augmented Graph Transformer (EGT): Edge channels influence attention by biasing query-key scores and gating values before aggregation, allowing structural information to control node-to-node information flow.The gated attention uses a clipped scaled dot product for numerical stability.
  • 3.2 Edge-augmented Graph Transformer (EGT): Global self-attention forms a dynamically weighted complete graph rather than following a fixed local neighborhood aggregation pattern.Random attention masking is used during training to encourage the network to exploit full connectivity.
  • 3.3 Dynamic Centrality Scalers: Centrality scalers multiply aggregated values using the logarithm of each node’s summed sigmoid gates, providing an adaptive centrality measure for dynamic attention graphs.The authors state that this makes global self-attention at least as powerful as the 1-WL isomorphism test and potentially more powerful through multiple-hop aggregation.
  • 3.4 SVD-based Positional Encodings: SVD-based positional encodings derive denoised edge information from top singular vectors and project it into the node-embedding dimension before addition to node inputs.The method randomly flips singular-vector signs during training to improve generalization; absolute coordinates may theoretically hamper generalization but did not do so in practice.

4 EXPERIMENTS AND RESULTS

EGT is evaluated across node-, edge-, and graph-level tasks, including large-scale molecular regression, ablations, transfer learning, and attention-pattern analyses. It generally outperforms local and convolutional GNNs, while global attention and evolving edge channels provide task-relevant advantages.

  • Medium-scale Performance: EGT outperforms other GNNs on all six benchmarking datasets except CIFAR10, where all models exhibit substantial overfitting.The comparison includes GAT, GT, and Graphormer under a fixed parameter-budget training setting.
  • Large-scale Performance: EGT achieves lower MAE than convolutional and local self-attention GNNs on PCQM4M and PCQM4Mv2, exceeding Graphormer on the molecular regression task.The reported results include small, medium, and large EGT models; PCQM4M results use the validation set because test labels are unavailable.
  • Large-scale Performance: EGT uses transfer learning from PCQM4Mv2 to fine-tune models on MolPCBA and MolHIV, selecting EGTLarger for transfer experiments.The selected 30-layer model achieves a validation MAE of 0.0869 on PCQM4M.
  • Ablation Study: EGT-Simple approaches EGT overall but is especially weaker on edge-related tasks and tasks where the distance objective cannot be applied, while both outperform EGT-Constrained on large PCQM4Mv2.EGT-Simple lacks dedicated residual edge channels and does not update edge embeddings across layers; EGT-Constrained restricts attention to 1-hop neighborhoods.
  • Ablation Study: Gates significantly improve performance, while attention-matrix dropout and dynamic centrality scalers support long-distance interactions and expressivity; virtual nodes and positional encodings yield more modest gains.These feature contributions are measured on the smallest EGT network using PCQM4Mv2.
  • Ablation Study: Attention patterns are selective and become increasingly non-local with depth on ZINC, while PCQM4Mv2 exhibits non-local aggregation from the lowest layers.The analysis compares adjacency, distances, individual and averaged attention heads, and hop weights across layers and datasets.

5 CONCLUSION AND FUTURE WORK

The paper introduces edge channels as a transformer extension for processing graph structure and concludes that local convolution is not essential for graph learning. Experiments support this conclusion across multiple scales and transfer-learning settings, including a new state-of-the-art result on PCQM4Mv2.

  • Edge channels extend transformers to accept, process, and output structural information such as links and edge labels.
  • Experimental results across medium-scale, large-scale, and transfer-learning settings show EGT outperforming convolutional GNNs.
  • EGT achieves a new state-of-the-art result on the large-scale PCQM4Mv2 molecular dataset.
  • The reported performance improvement is directly linked to the model’s non-local aggregation.
  • Future work targets transductive, semi-supervised, and unsupervised settings, plus sub-quadratic computation and memory costs.

A DATA AND CODE AVALABILITY

The paper uses publicly available datasets and provides code to reproduce its reported results.

  • All datasets used in the work are publicly available.
  • The medium-scale benchmarking datasets include PATTERN, CLUSTER, MNIST, CIFAR10, TSP, and ZINC.
  • The OGB-LSC PCQM4M and PCQM4Mv2 datasets and the OGB MolPCBA and MolHIV datasets are available through OGB.
  • Code for reproducing the reported results is available in the authors’ GitHub repository.

B.1 Medium-scale Experiments

Medium-scale experiments follow the benchmarking setup of Dwivedi et al. and tune architecture dimensions within fixed parameter budgets.

  • Experiments cover PATTERN, CLUSTER, MNIST, CIFAR10, TSP, and ZINC under the benchmarking setting suggested by Dwivedi et al.
  • Models use a specified parameter budget of either 100K or 500K.
  • The number of layers and the widths of node and edge channels are varied to optimize validation performance.
  • Training uses the Adam optimizer and reduces the learning rate by 0.5 when validation loss does not improve.

B.2 Large-scale Experiments

Large-scale training uses learning-rate scheduling, virtual nodes, attention dropout, and a fixed one-million-step training budget.

  • Large-scale PCQM4M and PCQM4Mv2 training requires learning-rate warmup followed by cosine decay.
  • Table 6 reports common hyperparameters for the medium-scale experiments on all datasets.
  • Virtual nodes provide a more scalable alternative to global average pooling by collecting additional graph-level information.
  • Attention dropout is used instead of random masking of attention matrices and is the only regularization method applied to all models.
  • All models are trained for a fixed number of 1 million gradient update steps.

B.3 Transfer Learning Experiments

The EGTLarger model was pre-trained on PCQM4Mv2 and fine-tuned on the OGB MolPCBA and MolHIV datasets using shortened training schedules and transfer-learning hyperparameters.

  • EGTLarger was fine-tuned on the OGB datasets MolPCBA and MolHIV after pre-training on PCQM4Mv2.The fine-tuning experiments used the same learning rate, warmup, and cosine-decay method as pre-training, but fewer total gradient update steps.
  • Each training and evaluation experiment was run four times with four different random seeds.

C DETAILS OF ABLATED VARIANTS

The ablated variants test the roles of evolving edge channels, global versus neighborhood-limited attention, and gating in EGT. They also describe the corresponding computational trade-offs and implementation choices.

  • EGT-Simple: EGT-Simple uses global self-attention without dedicated residual channels for updating pairwise edge information.Its input edge embeddings participate directly in aggregation and are not updated from layer to layer.
  • EGT-Simple: EGT-Simple performs edge classification from pairwise node embeddings and input edge features because its edge embeddings are not updated across layers.
  • EGT-Constrained: EGT-Constrained limits self-attention to each node’s 1-hop neighborhood and tracks edge embeddings only for existing neighboring edges.
  • Experimental settings: Table 7 lists dataset-specific medium-scale hyperparameters, while Tables 8 and 9 cover large-scale and transfer-learning experiments, respectively.
  • EGT-Constrained: EGT-Constrained can have sub-quadratic computational and memory costs depending on graph sparsity, but sparse GPU aggregation can be difficult to parallelize.Masked attention was used instead for faster training on datasets containing smaller graphs.
  • Ungated Variant: The ungated variant removes the value-gating mechanism while retaining the edge channels’ attention-bias role, testing the utility of gating.
  • EGT-Simple: In EGT-Simple, edge channels influence aggregation only through an attention bias derived from transformed input edge embeddings.
Loading 2108.03348v3…