Bayesian Sheaf Neural Networks

TL;DR

This paper introduces Bayesian sheaf neural networks that incorporate a variational approach to learn cellular sheaves, improving robustness and performance in heterophilic graph data tasks.

Contribution

It proposes a novel variational framework for learning cellular sheaves in graph neural networks, including a new family of distributions on $SO(n)$ using the Cayley transform.

Findings

Achieves leading performance on several graph datasets.

Less sensitive to hyperparameters with limited training data.

Demonstrates robustness of the Bayesian approach.

Abstract

Equipping graph neural networks with a convolution operation defined in terms of a cellular sheaf offers advantages for learning expressive representations of heterophilic graph data. The most flexible approach to constructing the sheaf is to learn it as part of the network as a function of the node features. However, this leaves the network potentially overly sensitive to the learned sheaf. As a counter-measure, we propose a variational approach to learning cellular sheaves within sheaf neural networks, yielding an architecture we refer to as a Bayesian sheaf neural network. As part of this work, we define a novel family of reparameterizable probability distributions on the rotation group $SO(n)$ using the Cayley transform. We evaluate the Bayesian sheaf neural network on several graph datasets, and show that our Bayesian sheaf models achieve leading performance compared to baseline models and are less sensitive to the choice of hyperparameters under limited training data settings.