Masked Attention is All You Need for Graphs
David Buterez Jon Paul Janet Dino Oglic Pietro Lio
Graph neural networks (GNNs) and variations of the message passing algorithm
are the predominant means for learning on graphs, largely due to their
flexibility, speed, and satisfactory performance. The design of powerful and
general purpose GNNs, however, requires significant research efforts and often
relies on handcrafted, carefully-chosen message passing operators. Motivated by
this, we propose a remarkably simple alternative for learning on graphs that
relies exclusively on attention. Graphs are represented as node or edge sets
and their connectivity is enforced by masking the attention weight matrix,
effectively creating custom attention patterns for each graph. Despite its
simplicity, masked attention for graphs (MAG) has state-of-the-art performance
on long-range tasks and outperforms strong message passing baselines and much
more involved attention-based methods on over 55 node and graph-level tasks. We
also show significantly better transfer learning capabilities compared to GNNs
and comparable or better time and memory scaling. MAG has sub-linear memory
scaling in the number of nodes or edges, enabling learning on dense graphs and
future-proofing the approach.