Learnable Commutative Monoids for Graph Neural Networks

TL;DR

This paper introduces learnable commutative monoids for graph neural networks, enabling efficient, parallelizable aggregation with competitive performance, addressing limitations of recurrent aggregators.

Contribution

It proposes a novel framework for constructing learnable, commutative, associative operators, leading to an aggregator with logarithmic depth and improved parallelism in GNNs.

Findings

LCM aggregator achieves performance comparable to recurrent aggregators.

The proposed method significantly improves parallelism and scalability.

Empirical results demonstrate the effectiveness of learnable commutative monoids.

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.