Transformers Generalize DeepSets and Can be Extended to Graphs & Hypergraphs

1 Introduction . Graph is a universal data modality used to model social networks [ 28 ] , chemical compounds [ 10 ] , biological structures [ 9 ] , and interactions in particle physics [ 21 , 29 ] . Recent graph neural networks ( GNNs ) adopt a message-passing scheme [ 39 , 37 , 32 ] , where the node features are propagated and aggregated recurrently according to the neighborhood structure given in the input graph . Despite the simplicity , the local and recurrent nature makes them unable to discover dependency between any two nodes with a distance longer than the message-passing steps [ 11 ] . The locality of message-passing is also known to be related to the over-smoothing problem that prevents scaling of GNNs [ 23 , 4 , 27 ] . An alternative approach is using a more general set of operations that involve global interactions while respecting the permutation symmetry of graphs . Such operations either produce the same output regardless of the node ordering of input graph ( permutation invariance ) , or commute with node reordering ( permutation equivariance ) . Message-passing is an equivariant operation , but is restricted to local neighborhoods defined by the input graph . Recently , Maron et . al . ( 2019 ) [ 26 ] characterized the full space of invariant and equivariant linear layers . It turned out that these layers span not only the local neighborhood interactions of message-passing GNNs , but also global interactions such as the one between disconnected nodes , and even edge-to-node , or node-to-edge interactions ( Figure 1 ) . Notably , their formulation naturally extends to layers with different input & output orders ( e.g. , edge in , node out ) , and higher-order layers for hypergraphs . In theory , an invariant MLP composed of these layers should be more powerful than message-passing GNNs as they can model long-range dependency between nodes ( Figure 1 ) . However , they are currently not widely adopted due to relatively low performance and high asymptotic memory complexity . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Contributions In this work , we present higher-order Transformers that address the low performance and high complexity of invariant MLPs . First , based on an observation that the renowned Transformer encoder generalizes first-order equivariant linear layers or DeepSets [ 35 ] , we formulate higher-order Transformer layers that generalize equivariant linear layers by extending self-attention to higher orders ( Figure 1 ) . Second , while Transformer layers with input and output orders k , l has O ( nk+l ) asymptotic complexity , we show that by leveraging the sparsity of input hypergraphs , we can obtain O ( m2 ) complexity given input with m hyperedges . Further adopting kernel attention approaches , we propose a variant that reduces the complexity to O ( m ) and theoretically prove that they are more expressive than message passing networks . Finally , we test higher-order Transformers on a range of tasks , and demonstrate that they achieve significant improvements over invariant MLPs and are highly competitive in performance and scalability to the state-of-the-art GNNs . 2 Preliminary . In this section , we first describe how ( hyper ) graphs can be treated as higher-order tensors . We then describe linear layers that operate on higher-order tensors while respecting node permutation symmetry . In particular , we analyze operations within linear layers for second-order tensors ( graphs ) and show they involve global interactions on top of local aggregation . Let us define some notations . We denote a set as { a , ... , b } , a tuple as ( a , ... , b ) , and denote [ n ] = { 1 , ... , n } . We denote the space of order-k tensors as Rnk×d where d is feature dimension . For an order-k tensor A ∈ Rnk×d , we use multi-index i = ( i1 , ... , ik ) ∈ [ n ] k to denote Ai = Ai1 , ... , ik ∈ Rd . Let Sn be the set of all permutations of [ n ] . π ∈ Sn acts on i by π ( i ) = ( π ( i1 ) , ... , π ( ik ) ) , and acts on an order-k tensor A by ( π ·A ) i = Aπ−1 ( i ) . ( Hyper ) graphs as tensors Generally , a ( hyper ) graph data G can be represented as a tuple ( V , A ) , where V is a set of n nodes and A ∈ Rnk×d encodes features attached to hyperedges . The type of the hypergraph is indicated by the order k of the tensor A. First-order tensor is a set of features ( e.g. , point cloud , bag-of-words ) where Ai is the feature of node i. Second-order tensor encodes edge features ( e.g. , adjacency ) where Ai1 , i2 is the feature of edge ( i1 , i2 ) . Generally , an order-k tensor encodes hyperedge features ( e.g. , mesh ) where Ai1 , ... , ik is the feature of hyperedge ( i1 , ... , ik ) . Permutation invariance and equivariance Our problem of interest is building a functional relation f ( A ) ≈ T between tensor A and target T . If T is a single output vector , we often require that f is permutation invariant , that it satisfies f ( π ·A ) = f ( A ) ; if T is a tensor T = T , we often require that f is permutation equivariant , satisfying f ( π ·A ) = π ·f ( A ) , for all π ∈ Sn and A ∈ Rn k×d . In a typical design setup where a neural network f is built using linear layers and non-linear activations , the construction of f reduces to finding invariant and equivariant linear layers . Invariant and equivariant linear layers We summarize invariant linear layers Lk→0 : Rn k×d → Rd′ and equivariant linear layers Lk→l : Rn k×d → Rnl×d′ identified in Maron et . al . ( 2019 ) [ 26 ] . Note that invariant layer is a special case of Lk→l with l = 0 . In summary , given an input A ∈ Rn k×d , the order-l output of an equivariant layer Lk→l can be written with i ∈ [ n ] k , j ∈ [ n ] l as : Lk→l ( A ) j = ∑ µ ∑ i Bµi , jAiwµ + ∑ λ Cλj bλ , ( 1 ) where wµ ∈ Rd×d ′ , bλ ∈ Rd ′ are weight and bias parameters , Bµ ∈ Rnk+l and Cλ ∈ Rnl are basis tensors ( will be defined ) , and µ and λ are equivalence classes of order- ( k + l ) and order-l multi-indices , respectively . The equivalence classes are defined upon equivalence relation ∼ that , for multi-indices i , j ∈ [ n ] k , i ∼ j iff ( i1 , ... , ik ) = ( π ( j1 ) , ... , π ( jk ) ) for some node permutation π ∈ Sn . Then , a multi-index i and all members j in its equivalence class have the same ( permutation-invariant ) equality pattern : ia = ib ⇔ ja = jb for all a , b ∈ [ k ] . Consequently , each equivalence class µ ( or λ ) is a distinct set of all order- ( k + l ) ( or order-l ) multi-indices having a specific equality pattern . Notably , we can represent each equivalence class of order-k multi-indices as a unique partition of [ k ] regardless of n , where the partition specifies the equality pattern . e.g. , with k = 2 , we have two partitions and respective equivalent classes : µ1 = { { i1 , i2 } } the set of all ( i1 , i2 ) with i1 = i2 , and µ2 = { { i1 } , { i2 } } the set of all ( i1 , i2 ) with i1 6= i2 . Thus , with b ( k ) the k-th Bell number or number of partitions of [ k ] , we have b ( k+ l ) equivalence classes µ for the weight and b ( l ) equivalence classes λ for the bias . We guide the readers interested in derivation to Maron et . al . ( 2019 ) [ 26 ] . With the equivalence classes , basis tensors are then defined as follows : Bµi , j = { 1 ( i , j ) ∈ µ 0 otherwise ; C λ j = { 1 j ∈ λ 0 otherwise ( 2 ) In Eq . ( 1 ) , each equivalence class µ determines which ( hyper ) edges participate and how they interact in summation ∑ i B µ i , jAi . As an example , let us consider L2→2 that maps input edges i = ( i1 , i2 ) to output edges j = ( j1 , j2 ) . An equivalence class µ = { { i1 , i2 } , { j1 , j2 } } represents all ( i , j ) ∈ µ that i1 = i2 , j1 = j2 , and i1 6= j1 . Then , due to masking by Bµ , only input elements Ai1 , i2 with i1 = i2 participate in the summation and gives output L2→2 ( A ) j1 , j2 for j1 = j2 , i1 6= j1 . Intuitively , this is analogous to computing a node feature by sum-pooling all the other nodes ( Fig . 2 ( b ) ) . Different µ accounts for other interactions as shown in Figure 2 . Notably , it contains a richer set of operations beyond local interactions modeled by message-passing ( Fig . 2 ( a ) ) , such as global interaction across all nodes ( Fig . 2 ( b ) ) , and edge-to-node ( Fig . 2 ( c ) ) , node-to-edge interactions ( Fig . 2 ( d ) ) . We finish the section by writing out the first-order equivariant layer L1→1 . As b ( 2 ) = 2 , the layer has two equivalence classes µ1 = { { i1 , j1 } } and µ2 = { { i1 } , { j1 } } . Then , we have Bµ1 = In and Bµ2 = 1n1 > n − In , with 1n ∈ Rn vector of ones . Then , given a set of features A ∈ Rn×d , L1→1 ( A ) = InAw ′ 1 + ( 1n1 > n − In ) Aw′2 + 1nb > ( 3 ) = InAw1 + 1n1 > nAw2 + 1nb > , ( 4 ) where w1 , w2 , w′1 , w ′ 2 ∈ Rd×d ′ , b ∈ Rd′ . L1→1 is analogous to a combination of elementwise feedforward ( µ1 ) and sum-pooling of set elements ( µ2 ) , and is also known as a DeepSet layer [ 35 ] . 3 Higher-Order Transformers . In Section 2 , we introduced higher-order linear equivariant layers Lk→l , and showed that they contain various global and node/edge interactions that are not covered by message-passing . In this section , we establish a connection between the first-order equivariant layer L1→1 and self-attention of Transformer encoder layers [ 31 ] . Then , we extend the relationship to higher orders by tensorizing queries and keys , and formulate higher-order Transformer layers . We finish the section by proposing a principled parameter reduction for queries and keys , which reduces a fair amount of computation .

Summary: In this paper, the authors proposed a transformer layer that generalizes the equivariant linear layers with self-attention to higher order tensor input. They defined the attention coefficient for each equivalence class $\mu$ with indices satisfied $(i, j) \in \mu$ to reduce unnecessary computation between queries and keys matrices. To further reduce the computation complexity of higher order transformers, the authors provides 3 steps of approximation: i) a lightweight linear layers defined only for subset of equivalence class to decouple the indices $i$ and $j$. ii) adding proxy hyperedges $E'$ that is contained in original hyperedges $E$ and help reducing the computation to only $m$ hyperedges across different layers. iii) a kernel trick to decompose the attention coefficients into kernel map of query and keyword matrices.