Autobahn: Automorphism-based Graph Neural Nets

1 Introduction . The successes of artificial neural networks in domains such as computer vision and natural language processing have inspired substantial interest in developing neural architectures on graphs . Since graphs naturally capture relational information , graph-structured data appears in a myriad of fields . However , working with graph data raises new problems . Chief among them is the problem of graph isomorphism : for the output of our neural network to be reliable , it is critical that the network gives the same result independent of trivial changes in graph representation such as permutation of nodes . Considerable effort has gone into constructing neural network architectures that obey this constraint [ 1– 5 ] . Arguably , the most popular approach is to construct Message-Passing Neural Networks ( MPNNs ) [ 6 , 7 ] . In each layer of an MPNN , every node aggregates the activations of its neighbors in a permutation invariant manner and applies a linear mixing and nonlinearity to the resulting vector . While subsequent architectures have built on this paradigm , e.g. , by improving activations on nodes and graph edges [ 8 , 9 ] , the core paradigm of repeatedly pooling information from neighboring nodes has remained . These architectures are memory efficient , intuitively appealing , and respect the graph ’ s symmetry under permutation of its nodes . However , practical results have shown that they can oversmooth signals [ 10 ] and theoretical work has shown that they have trouble distinguishing certain graphs and counting substructures [ 8 , 11–13 ] . Moreover , MPNNs do not use the graph ’ s topology to its fullest extent . For instance , applying an MPNN to highly structured graphs such as grid graphs does not recover powerful known architectures such as convolutional neural networks . It is also not clear how to best adapt MPNNs to families of graphs with radically different topologies : MPNNs for citation graphs and molecular graphs are constructed in largely the same way . This suggests that it should be possible to construct more expressive graph neural networks by directly leveraging the graph ’ s structure . In this work , we introduce a new framework for constructing graph neural networks , Automorphismbased Neural Networks ( Autobahn ) . Our research is motivated by our goal of designing neural 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . networks that can learn the properties of small organic molecules accurately enough to make a significant contribution to drug discovery and materials design [ 14–17 ] . The properties of these molecules depend crucially on multi-atom substructures , making the difficulties MPNNs have in recognizing substructures a critical problem . We realized that we could circumvent this limitation by making graph substructures themselves the fundamental units of computation . The symmetries of our substructures then inform the computation . For instance , the benzene molecule forms a ring of six atoms and is a common subunit in larger molecules . On this ring we have a very natural and mathematically rigorous notion of convolution : convolution on a one-dimensional , periodic domain . This symmetry is encoded by a graph ’ s automorphism group : the group that reflects our substructures ’ internal symmetries . Our networks directly leverage the automorphism group of subgraphs to construct flexible , efficient neurons . Message passing neural networks arise naturally in this framework when the substructures used are local star graphs , and applying Autobahn to grid graphs can recover standard convolutional architectures such as steerable CNNs . More generally , Autobahn gives practitioners the tools to build bespoke graph neural networks whose substructures reflect their domain knowledge . As an example , in Section 6 , we present a novel architecture outside of the message-passing paradigm that nevertheless achieves performance competitive with state-of-the-art MPNNs on molecular learning tasks . 2 Graph Neural Networks . Neural Networks operate by composing several learned mappings known as “ layers ” . Denoting the ` ’ th layer in the network as φ ` , the functional form of a neural network Φ can be written as Φ = φL ◦ φL−1 ◦ . . . ◦ φ1 . Each layer is typically constructed from a collection of parts , the titular “ neurons ” . We denote the i ’ th neuron in the ` ’ th layer as n ` i , and denote its output ( the “ activation ” of the neuron ) as f ` i . Architectures differ primarily in how the neurons are constructed and how their inputs and outputs are combined . When constructing a neural network that operates on graph data , care must be taken to preserve the input graphs ’ natural symmetries under permutation . Let G be a graph with node set { v1 , . . . , vn } , adjacency matrix A∈Rn×n and d-dimensional node labels bi ∈ Rd stacked into a matrix B ∈Rn×d . Permuting the numbering of the nodes of G by some permutation σ : { 1 , 2 , . . . , n } → { 1 , 2 , . . . , n } transforms A 7→ Aσ Aσi , j = Aσ−1 ( i ) , σ−1 ( j ) , ( 1 ) and B 7→ Bσ Bσi , n = Bσ−1 ( i ) , n. ( 2 ) This transformation does not change the actual topology of G. Consequently , a fundamental requirement on graph neural networks is that they be invariant with respect to such permutations . 2.1 Message-Passing Neural Networks . Message-passing neural networks ( MPNNs ) have emerged as the dominant paradigm for constructing neural networks on graphs . Every neuron in an MPNN corresponds to a single node in the graph . Neurons aggregate features from neighboring nodes by passing them through a permutation invariant function . They then combine the result with the node ’ s original message , and pass the result through a learned linear function and nonlinearity ( Figure 1 ) [ 8 ] . Since information is transmitted over the graph ’ s topology , MPNNs are not confounded by permutations of the graph . Specific architectures may differ in the details of the aggregation function [ 6 , 8 , 18 , 19 ] , may include additional terms to account for edge features [ 8 , 9 ] , may use complex transformations to construct the node features encoding local structure [ 20 , 21 ] , or may augment the graph with additional nodes [ 22 , 23 ] . 3 Permutation Equivariance . As discussed in the Introduction , MPNNs have fundamental limits to their expressiveness . To construct more powerful neural networks , we turn to the general formalism of group equivariant networks [ 24–27 ] . Our desire that permutations of the input graph leave our network ’ s output unaffected is formalized by the notion of group-invariance . Let us assume that our input data lives in a space X that is acted on by a group G , and for all g ∈ G denote the associated group action on X by Tg . The invariance constraint amounts to requiring : Φ = Φ ◦ Tg ∀g ∈ G. ( 3 ) One way to satisfy this would be to require that each layer φ ` be fully invariant to G. However , in practice this condition can be extremely restrictive . For this reason , networks commonly use group equivariant layers . Let X and Y be the input and output spaces of φ ` , with group actions Tg and T ′g , respectively . We say φ ` is equivariant to G if it obeys T ′g ◦ φ ` = φ ` ◦ Tg ∀g ∈ G. ( 4 ) This condition is weaker than invariance : we recover invariance when T ′g maps every element of Y to itself . Moreover , it is simple to show that the composition of two equivariant layers is also equivariant . Consequently , in all layers but the last we can enforce the weaker condition of equivariance and merely enforce invariance in the final layer . In the case of graph neural networks , the relevant group is the group of permutations : Sn ( called the symmetric group of degree n ) . 3.1 Permutation-equivariant networks . Recent work has developed a generic recipe for constructing group equivariant networks . In this formalism , any object that transforms under a group action is treated as a function on the group [ 24 ] ( see Section 1 in the supplement for a brief review ) . This allows all layers equivariant to the group to be treated using the same formalism , independently of how inputs and outputs transform . In particular , it can be shown that the only group-equivariant linear operation possible is a generalized notion of group convolution . For discrete groups , this convolution can be written as : ( f ∗ w ) ( u ) = ∑ v∈G f ( uv−1 ) w ( v ) . ( 5 ) To construct an equivariant neuron , we apply ( 5 ) to convolve our input activation f ` −1 with a learned weight function w , add a bias , and then apply a fixed equivariant nonlinearity , ν . Applying this approach to specific groups recovers the standard convolutional layers used in convolutional neural networks ( CNNs ) . For instance , applying ( 5 ) to the cyclic group of order n gives ( f ∗ w ) j = n∑ k=0 f ( rj−k ) w ( rk ) , ( 6 ) where r is the group element corresponding to rotation by 360/n degrees . Similarly , applying ( 5 ) to one-dimensional or two-dimensional discrete translation groups recovers the standard convolutions used for image processing . Instantiating this theory with the symmetric group has been successfully used to construct permutation equivariant networks for learning tasks defined on sets [ 28 ] and for graph neural networks [ 29–32 ] . However , enforcing equivariance to all permutations can be very restrictive . As an example , consider a layer whose domain and range transform according to ( 2 ) with cin incoming and cout outgoing channels . Let w1 , w2 ∈ Rcout×cin be learned weight matrices , and fj ∈ Rcin be the vector of incoming activations corresponding to node j . The most general possible convolution is then [ 28 ] ( f ∗ w ) i = w1fi + w2 n∑ j=1 fj , ( 7 ) but this is quite a weak model as activations from different nodes only interact through their sum . Consequently , relationships between nodes can only be captured in aggregate . Algorithm 1 Automorphism-based Neuron Input : f ` −1 . Incoming activation associated with G AG . Adjacency matrix of G AT . Adjacency matrix of the saved template graph . 1 : Find µ ∈ Sn such that AµG = AT . 2 : ( f ` −1 ) µ ← Tµ ( f ` −1 ) . Apply µ to incoming activation . 3 : ( f ` ) µ ← ν ( ( f ` −1 ) µ ∗ w + b ) . Convolution is over Aut ( T ) . 4 : f ` ← T ′µ−1 ( ( f ` ) µ ) . Map output to original ordering Return : f ` To address this fundamental limitation , several recent works considered improving the expressivity of MPNNs by defining higher order activations corresponding to pairs , triplets , or , in general , k–tuples of nodes [ 29 , 32 , 33 ] . Mathematically , this requires considering not just ( 1 ) and ( 2 ) , but the action of the symmetric group on k ’ th order tensors , Aσi1 , ... , ik = Aσ−1 ( i1 ) , ... , σ−1 ( ik ) . However , this can be prohibitively expensive for many nodes . For instance , organic chemistry depends crucially on the existence of aromatic rings , typically of six atoms . Manipulating sixth order tensors would be extremely costly .

Summary: This paper introduces Automorphism-based Graph Neural Networks (Autobahns), which decompose the graph into a collection of subgraphs and apply convolutions that are equivariant to each subgraph’s automorphism group. The specific architecture is determined by the choices of local neighborhoods and subgraphs. The proposed formalism not only generalizes the message passing GNNs but also encompasses novel architectures. Experiments are conducted on molecular graph datasets and demonstrate that Autobahns can achieve state-of-the-art results.