Relaxing Permutation Equivariance in Graph Neural Networks

Graph Neural Networks (GNNs) have been established as powerful methods for learning functions on graph-structured data, in which nodes carry no canonical ordering. Permutation equivariance to node index relabelling is therefore a defining inductive bias of GNNs. We argue that the representation of nodes in sets is limiting. We propose an alternative approach that represents and learns on observed graphs, in which nodes are augmented with one or several partial orders. Leveraging the equivalence between partially ordered sets and Directed Acyclic Graphs (DAGs), we introduce a new data structure called Partially Ordered Graphs (POGs). Building on theoretical work defining convolution operators on DAGs, we construct the Partially Ordered GNN (POGNN) architecture which combines trainable convolutions on both the node partially ordered sets and the underlying observed graph. We also prove that our POGNNs relax the permutation equivariance assumption of traditional GNNs in a tunable manner by being equivariant to the automorphism of the poset. The choice of partial order allows us to cover the full spectrum from node sets to node sequences. In practice, we demonstrate that our model improves the performance of traditional GNN approaches on several graph-level biochemical benchmark datasets.