Understanding Attention and Generalization in Graph Neural Networks

TL;DR

This paper investigates the role of attention mechanisms in graph neural networks, revealing conditions under which attention improves performance and proposing a weakly-supervised training method to enhance generalization on various datasets.

Contribution

It provides a controlled study of attention in GNNs, identifies when attention is beneficial, and introduces a weakly-supervised training approach to improve attention effectiveness.

Findings

Attention can be harmful or negligible under typical conditions.

Properly trained attention can boost classification performance by over 60%.

Weakly-supervised training of attention approaches supervised performance and outperforms unsupervised methods.

Abstract

We aim to better understand attention over nodes in graph neural networks (GNNs) and identify factors influencing its effectiveness. We particularly focus on the ability of attention GNNs to generalize to larger, more complex or noisy graphs. Motivated by insights from the work on Graph Isomorphism Networks, we design simple graph reasoning tasks that allow us to study attention in a controlled environment. We find that under typical conditions the effect of attention is negligible or even harmful, but under certain conditions it provides an exceptional gain in performance of more than 60% in some of our classification tasks. Satisfying these conditions in practice is challenging and often requires optimal initialization or supervised training of attention. We propose an alternative recipe and train attention in a weakly-supervised fashion that approaches the performance of supervised models, and, compared to unsupervised models, improves results on several synthetic as well as real datasets. Source code and datasets are available at https://github.com/bknyaz/graph_attention_pool.