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 proposed Autobahn, a new GNN model for learning representations of graphs for graph classification tasks. The idea is to extract subgraphs that are isomorphic to predefined graph templates (e.g., cycles, paths) and perform a convolution operation on its automorphism group. This paper proposed two operations to perform interaction between the representations on different subgraphs: narrowing and promoting. This paper applied the proposed model to graph classification tasks on a molecular graph dataset and verified its effectiveness.