Learnable Commutative Monoids for Graph Neural Networks

Euan Ong, Petar Veličković

Keywords: graph neural networks, aggregation, commutative monoids, principal neighbourhood aggregation

Abstract: Graph neural networks (GNNs) have been shown to be highly sensitive to the choice of aggregation function. While summing over a node's neighbours can approximate any permutation-invariant function over discrete inputs, Cohen-Karlik et al. [2020] proved there are set-aggregation problems for which summing cannot generalise to unbounded inputs, proposing recurrent neural networks regularised towards permutation-invariance as a more expressive aggregator. We show that these results carry over to the graph domain: GNNs equipped with recurrent aggregators are competitive with state-of-the-art permutation-invariant aggregators, on both synthetic benchmarks and real-world problems. However, despite the benefits of recurrent aggregators, their O(V) depth makes them both difficult to parallelise and harder to train on large graphs. Inspired by the observation that a well-behaved aggregator for a GNN is a commutative monoid over its latent space, we propose a framework for constructing learnable, commutative, associative binary operators. And with this, we construct an aggregator of O(log V) depth, yielding exponential improvements for both parallelism and dependency length while achieving performance competitive with recurrent aggregators. Based on our empirical observations, our proposed learnable commutative monoid (LCM) aggregator represents a favourable tradeoff between efficient and expressive aggregators.

Type Of Submission: Full paper proceedings track submission (max 9 main pages).

Agreement: Check this if you are okay with being contacted to participate in an anonymous survey.

TL;DR: We leverage a perspective of learnable commutative monoids to propose an expressive, flexible and efficient aggregator for GNNs.

Type Of Submission: Full paper proceedings track submission.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2212.08541/code)