Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs
Year: 2022

Abstract:
We present a generalization of Transformers to any-order permutation invariant data (sets, graphs, and hypergraphs). We begin by observing that Transformers generalize DeepSets, or first-order (set-input) permutation invariant MLPs. Then, based on recently characterized higher-order invariant MLPs, we extend the concept of self-attention to higher orders and propose higher-order Transformers for order-k data (k = 2 for graphs and k > 2 for hypergraphs). Unfortunately, higher-order Transformers turn out to have prohibitive complexityO(n) to the number of input nodes n. To address this problem, we present sparse higher-order Transformers that have quadratic complexity to the number of input hyperedges, and further adopt the kernel attention approach to reduce the complexity to linear. In particular, we show that the sparse second-order Transformers with kernel attention are theoretically more expressive than message passing operations while having an asymptotically identical complexity. Our models achieve significant performance improvement over invariant MLPs and message-passing graph neural networks in large-scale graph regression and set-to-(hyper)graph prediction tasks. Our implementation is available at https://github.com/jw9730/hot.

heading: 1 Introduction
text: 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.

heading: 2 Preliminary
text: 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].

heading: 3 Higher-Order Transformers
text: 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.

heading: 3.1 Transformers generalize DeepSets
text: As shown in Section 2, first-order linear layer, or DeepSet, has a simple structure composed of feedforward and sum-pooling (Eq. (4)). Although it is theoretically proven to be a universal approximator of permutation-invariant functions [35], static sum-pooling could be limited in capturing interactions of set elements, motivating the use of sophisticated pooling. In particular, the self-attention mechanism of Transformer encoder [31] was shown to achieve a large performance gain in set modeling via context-aware weighted pooling [22, 34]. To see this, let us write out the Transformer encoder layers.
A Transformer encoder layer is a function Enc : Rn×d → Rn×d consisting of two layers: a selfattention layer Attn : Rn×d → Rn×d and an elementwise feedforward layer MLP : Rn×d → Rn×d. For a set of n input vectors X ∈ Rn×d, a Transformer layer computes the following:
Attn(X)i = Xi + H∑ h=1 n∑ j=1 αhijXjw V h w O h , (5)
Enc(X)i = Attn(X)i + MLP(Attn(X))i, (6)
where αh = σ(XwQh (Xw K h ) >) is an attention coefficient with an activation σ, H is the number of heads, dH is head size, dF is hidden dimension, and wOh ∈ RdH×d, wVh , wKh , w Q h ∈ Rd×dH .1
Now, we show that Transformer layers are generalized first-order linear equivariant functions. By setting αhij = 1 and assuming that MLP(Y ) approximates a linear layer Y w
F + bF> following the universal approximation theorem [16], the Transformer layer reduces to the following 2 :
Enc(X)i = Xi(In + wF ) + n∑ j=1 Xjw V O(In + w F ) + bF>, (7)
where wV O = ∑H h=1 w V h w O h . This is equivalent to a DeepSet layer in Eq. (4) with w1 = In + w
F , w2 = w V O(In + w F ), and b = bF . In other words, we can convert DeepSets to Transformers by changing the static pooling to attention and replacing elementwise linear mapping with an MLP. We generalize this approach to higher-order input and output, which leads to the formulation of higher-order Transformers for graphs and hypergraphs.

heading: 3.2 Higher-order Transformer layers
text: In Section 3.1, we showed that Transformer layers are generalized first-order equivariant linear layers L1→1. Notably, the generalization procedure was equivalent to changing static pooling to attention and adding feedforward MLP. In this section, we generalize the approach to Lk→l with arbitrary orders k and l and formulate higher-order Transformer layer Enck→l.
In general, we define higher-order Transformer layer as a function Enck→l : Rn k×d → Rnl×d with two layers: a higher-order self-attention layer Attnk→l : Rn k×d → Rnl×d and a feedforward layer MLPl→l : Rn l×d → Rnl×d. For an input tensor A ∈ Rnk×d, a Transformer layer computes:
MLPl→l(Attnk→l(A)) = L2l→l ( ReLU(L1l→l(Attnk→l(A))) ) , (8)
Enck→l(A) = Attnk→l(A) + MLPl→l(Attnk→l(A)), (9)
where L1l→l : Rn l×d→nl×dF and L2l→l : Rn l×dF→nl×d are equivariant linear layers with hidden dimension dF . Remaining question is how to define and compute higher-order self-attention Attnk→l(A).
To design Attnk→l, we remove the bias from Eq. (1) and introduce attention coefficients. It is achieved by changing each Bµ ∈ Rnk+l to an attention coefficient tensor αh,µ ∈ Rnk+l with multiple heads:
Attnk→l(A)j = H∑ h=1 ∑ µ ∑ i αh,µi,j Aiw V h,µw O h,µ, (10)
where wOh,µ ∈ RdH×d, wVh,µ ∈ Rd×dH are learnable parameters, H denotes the number of heads, and dH denotes head size. Then similar to first-order case (Section 3.1), we can show the following:
1Note that we omitted normalization after Attn(·) and MLP(·) for simplicity as in [34, 12]. 2In practice, Transformer employs softmax in attention and deviates from Deepsets due to normalization.
Theorem 1. Enck→l (Eq. (9)) is a generalization of Lk→l (Eq. (1)).
Proof. Let αh,µi,j = 1 for all h, µ, and (i, j) ∈ µ. This renders α h,µ = Bµ from the definition of Bµ. Additionally, let MLPl→l(Attnk→l(A))j = ∑ λC λ j bλ for all j ∈ [n]l. That is, MLPl→l ignores input and reduces to an invariant bias in Eq. (1). Then, Eq. (9) reduces to the following:
Enck→l(A)j = ∑ µ ∑ i Bµi,jAi H∑ h=1 wVh,µw O h,µ + ∑ λ Cλj bλ, (11)
which is equivalent to Eq. (1) with wµ = ∑H h=1 w V h,µw O h,µ.
Now, we describe how to compute each attention tensor αµ ∈ Rnk+l from input A ∈ Rnk×d (Eq. (10), we drop head index h for brevity). We obtain each attention tensor from higher-order query and key:
αµi,j =
{ σ(Qµj ,K µ i )/Zj (i, j) ∈ µ
0 otherwise where Qµ = Lµk→l(A),K µ = Lµk→k(A). (12)
where Zj = ∑ i|(i,j)∈µ σ(Q µ j ,K µ i ) is a normalization constant. Note that query and key tensors are computed from the input A using the equivariant linear layers in Eq. (1). This leads to permutation equivariance (or invariance) of Transformer encoder layer Enck→l in Eq. (9).
Reducing orders of query and key Although Eq. (10) and Eq. (12) provide a simple and generic definition of higher-order self-attention, we observe that there exist a lot of unnecessary computations. Specifically, there exist elements of query Qµ and key Kµ that are unused in computation of attention coefficient αµi,j, as it depends only on indices satisfying (i, j) ∈ µ. In fact, it turns out that the effective orders of query and key are much smaller than l and k, as we show below:
Proposition 1. From Eq. (12), let u(·) denote the number of unique entries in a multi-index. With uq = u(j), uk = u(i) for some (i, j) ∈ µ, we can always find suitable linear layers Lµk→uq , L µ k→uk and index space mappings fµq : [n] l → [n]uq , fµk : [n]k → [n]uk that satisfy the following.
αµi,j = σ(Q̃ µ j′ , K̃ µ i′)/Z̃j ∀(i, j) ∈ µ, (13)
where Z̃j = ∑ i|(i,j)∈µ σ(Q̃ µ j′ , K̃ µ i′), Q̃ µ = Lµk→uq (A), K̃ µ = Lµk→uk(A), j ′ = fµq (j), i ′ = fµk (i).
Proof. We leave the proof in Appendix A.1.1.
Based on Proposition 1, we can compute query and key in Eq. (12) in a much compact way using linear layers with output orders uq and uk instead of l and k. In our experiments, we observe that this optimization is very useful in reducing the number of parameters and memory footprint to a feasible level without affecting the effective model capacity.

heading: 4 Asymptotically Efficient Higher-Order Transformers
text: In Section 3, we formulated higher-order Transformer layers Enck→l and showed that they generalize linear equivariant layers Lk→l. However, this capability comes with a cost; the high asymptotic complexity of the Transformer encoder limits its practical merits. Specifically, we show the following:
Property 1. Given input size n, the asymptotic complexity of a linear layer Lk→l (Eq. (1)) is O(nk+l), and complexity of an encoder layer Enck→l (Eq. (9)) is O(nk+l + n2k + n2l).
Proof. We leave the proof in Appendix A.1.2.
Thus, in this section, we further analyze the encoder layer and propose a number of optimization and relaxation to reduce the asymptotic complexity with a minimal impact on capability. Notably, combining all our strategies reduces the complexity to O(m) given a hypergraph with m hyperedges. Even with such efficiency, we show that the reduced version of higher-order Transformer is theoretically more expressive than all message-passing neural networks.

heading: 4.1 Linear layers with reduced complexity
text: A major computation bottleneck within Transformer encoder layer Enck→l is the higher-order linear layer, since it is used for key, query, and feedforward layer. By exploiting only a subset of basis, we show that we can implement lightweight version of linear layer with reduced asymptotic complexity. Proposition 2. Given a linear layer Lk→l in Eq. (1), we can always find a nonempty subsetM of equivalence classes such that computation of the following for all j has O(nl) complexity.
L̄k→l(A)j = ∑ µ∈M ∑ i Bµi,jAiwµ + ∑ λ Cλj bλ, (14)
Proof. We leave the proof in Appendix A.1.3. We term the reduced linear layer L̄k→l in Eq. (14) as a lightweight linear layer. In practice, we chooseM among equivalence classes that do not involve summation over input. For instance, for L̄1→1, we use only the basis for elementwise mapping (In) and drop sum-pooling (1n1>n ) from Eq. (4). This approximation effectively reduces the complexity of linear layers at the cost of losing inter-element dependencies. We employ the lightweight linear layers within the Enck→l (Eq. (9)) to compute query and key embeddings, while the element dependency within Enck→l is handled by the higher-order self-attention using all equivalence classes as in Eq. (12). This design is coherent to original (first-order) Transformers, where the elements are first linearly projected to query/key with elementwise basis (In) and interaction is handled by attention (implicitly 1n1>n ). Importantly, this does not hurt Theorem 1 as attention coefficients can still reduce to one and MLP can reduce to bias. By using lightweight linear layers in Enck→l, we can significantly reduce the computational cost. The complexity of Lµk→uk , L µ k→uq , and MLPl→l in Eq. (13) and Eq. (9) reduces to O(n
k), O(nl), and O(nl), respectively, and as a result Enck→l becomes O(nk+l). As a result, we have higher-order Transformer layers Enck→l that generalize Lk→l while retaining the complexityO(nk+l). From here we assume that all linear layers for key, query, and MLP are lightweight.

heading: 4.2 Sparse Transformer layers
text: Even with lightweight linear layers, O(nk+l) complexity of Enck→l is still impractical. For example, for graphs, complexity larger thanO(n2) is regarded prohibitive while Enc2→2 isO(n4). Fortunately, leveraging the sparsity inherent in real-world graphs can significantly reduce the complexity. As we show, it reduces the complexity of Enck→l toO(m2) for processing a hypergraph with m hyperedges. Let E the set of hyperedges of input hypergraph G. Each hyperedge is generally represented by a multi-index i ∈ [n]k, so we denote E = {i1, ..., im} with m hyperedges. Leveraging sparsity of E is straightforward when the order of the hypergraph is fixed (e.g., Lk→k); we can perform computations with only respect to the existing hyperedges i, j ∈ E. However, in our framework, order can change by layers (e.g., Lk→l), making it difficult to directly transfer the sparsity structure E to different-order output. Inspired by network projection [5], we remedy this by constructing E′ such that for any j ∈ E′, there exists i ∈ E containing all unique elements of j. For example, for L3→2, this corresponds to projection of third-order E to second-order E′ by obtaining edges from sides of triangles. Despite simplicity, this simple heuristic works well in general and generalizes to any order.
Then, we integrate the hyperedge sets E, E′ into computation of linear layer in Eq. (14) as follows:
L̄k→l(A, E)j =
{ ∑ µ∈M ∑ i∈E B µ i,jAiwµ + ∑ λC λ j bλ j ∈ E′
0 otherwise (15)
Likewise, integrating E, E′ into attention computation in Eq. (10) and Eq. (12) we have: Attnk→l(A, E)j = { ∑H h=1 ∑ µ ∑ i∈E α h,µ i,j Aiw V h,µw O h,µ j ∈ E′
0 otherwise (16) where αh,µi,j = { σ(Qh,µj ,K h,µ i )/Zj (i, j) ∈ µ, i ∈ E, j ∈ E′
0 otherwise , (17)
where Qh,µ = L̄h,µk→l(A, E), K h,µ = L̄h,µk→k(A, E), Zj = ∑ i|(i,j)∈µ∧i∈E σ(Q h,µ j ,K h,µ i ).
With the computations, we can show the following: Property 2. When given E with m elements, the equivariant linear layer in Eq. (15) has O(m) complexity, and self-attention computation in Eq. (16) has O(m2) complexity. Consequently, the computation done by a Enck→l composed of the layers has O(m2) complexity. Proof. We leave the proof in Appendix A.1.5.

heading: 4.3 Kernel attention trick
text: Section 4.2 shows that, by constraining linear layers within Enck→l to have sparse input and output, we can reduce O(nk+l) complexity to quadratic O(m2) to input size. Yet, even O(m2) can be demanding with large or dense input. As the quadratic term comes from self-attention computation, we follow the prior work in kernel attention [18, 6] and view attention coefficients as pairwise dot-product scores. As we will show, this allows us to further reduce the complexity of Enck→l to O(nk + nl) and even O(m) for sparse version, at the cost of relaxing some modeling assumption. We begin by approximating attention coefficient in Eq. (12) using pairwise dot-product kernel [18, 6]:
αµi,j =
{ φ(Qµj )
>φ(Kµi )/Zj (i, j) ∈ µ 0 otherwise
where Zj = ∑
i|(i,j)∈µ
φ(Qµj ) >φ(Kµi ) (18)
where φ : RdH → RdK+ is kernel feature map. The choice of kernel can be flexible, and in our implementation we adopt Performer kernel [6] that has strong theoretical and empirical guarantee.
Substituting Eq. (18) in Eq. (10), we have: Attnk→l(A)j = ∑ µ Z−1j ∑ i|(i,j)∈µ φ(Q µ j ) >φ(Kµi )Aiw V µ w O µ
= ∑ µ Z−1j φ(Q µ j ) > ∑ i|(i,j)∈µ φ(Kµi )Aiw V µ w O µ , (19)
where Zj = φ(Q µ j ) > ∑
i|(i,j)∈µ
φ(Kµi ). (20)
Here, the inner-summations in Eq. (19) and Eq. (20) are over the key index i, which is coupled with the query index j by (i, j) ∈ µ. This coupling causes the major computational bottleneck, since the inner-summations should be computed for every j-th query. We propose an approximation by decoupling the key and query indices and taking inner-summations over I = ⋃ j {i|(i, j) ∈ µ}:
Attnk→l(A)j ≈ ∑ µ Z−1j φ(Q µ j ) > ∑ i∈I φ(Kµi )Aiw V µ w O µ , where Zj ≈ φ(Q µ j ) > ∑ i∈I φ(Kµi ). (21)
The approximation allows a query to attend to some additional keys (depending on µ), but it does not hurt the central requirement of attention that each query can assign different attention weights to the keys. With the approximation, we can compute the summations ∑ i∈I φ(K µ i )Ai and ∑ i∈I φ(K µ i ) only once and reuse them across all query indices j to reduce the cost. Specifically, we show: Property 3. The encoder Enck→l with approximation in Eq. (21) has a complexity of O(nk + nl). Exploiting sparsity further reduces the complexity to O(m), linear to the number of hyperedges m. Proof. We leave the proof in Appendix A.1.6.

heading: 4.4 Theoretical analysis and comparison to message-passing
text: We showed that exploiting sparsity and adopting kernel attention reduces the computational complexity of the Transformer encoder to linear to input edges. Considering graphs (k = l = 2), this complexity is equivalent to or better than O(n+m) complexity of the message passing operation3. Still, we show that our (approximate) encoder is theoretically more expressive than message passing.
Specifically, we show the following: Theorem 2. A composition of two sparse Transformer layers Enc2→2 with kernel attention can approximate any message passing algorithms (Gilmer et. al. (2017) [10]) to arbitrary precision, while the opposite is not true. Proof. We leave the proof in Appendix A.1.7.
This leads to the following corollary: Corollary 1. Second-order sparse Transformers with kernel attention are more expressive than any message-passing neural networks within the framework of Gilmer et. al. (2017) [10].
In the proof, we note that local information propagation in message-passing GNNs is carried out in Enc2→2 by a single µ = {{i1}, {i2, i3, i4}}. Other types of µ would carry out different operations, which provides intuition on the powerfulness of Transformers.
3In practice, we place node features on the diagonals of the adjacency matrix, leading to O(n+m).

heading: 5 Experiments
text: In this section, we demonstrate the capability of higher-order Transformers on a variety of tasks including synthetic data, large-scale graph regression, and set-to-(hyper)graph prediction. Specifically, we use a synthetic node classification dataset from Gu et. al. (2020) [11], a molecular graph regression dataset from Hu et. al. (2021) [17], two set-to-graph prediction datasets from Serviansky et. al. (2020) [29], and three hyperedge prediction datasets used in Zhang et. al. (2020) [36]. Details including the datasets and hyperparameters can be found in Appendix A.2. We implemented invariant MLP in Maron et. al. (2019) [26] as one of baselines, which we abbreviate as MLPπ. We build our model by gradually adding lightweight linear layers (Sec. 4.1), sparse linear layers (Sec. 4.2), and kernel attention (Sec. 4.3) and denote them by (D), (S), and (φ), respectively.
Runtime and memory analysis To experimentally verify our claims on linear complexity in Sec. 4.3, we conducted a runtime and memory consumption analysis on all second-order models using random graphs; details can be found in Appendix A.2.3. The results are shown in Figure 3. Consistent with our theoretical claims, sparse second-order Transformer with kernel attention (Ours (S, φ)) is the only variant that linearly scales to input size in terms of both time and memory. Also, it is the only one that successfully scales to graphs with > 50k edges, while still very fast in smaller graphs.
Synthetic chains To test the ability of higher-order Transformers in modeling long-range interactions in graphs, we used a synthetic dataset where the task is node classification in chain graphs. Binary class information is provided in a terminal node, so a model is required to propagate the information across the chain by handling long-range dependency. We used training and test sets with chains of length 20 and 200 respectively. As baselines, we used 3 message-passing networks, a sparse invariant MLP, and ablated versions of sparse second-order Transformers where µ for global pooling are removed (w/o global). Further details can be found in Appendix A.2.4.
The test performances are in Table 1. We first note that second-order Transformers successfully capture long-range dependency up to 200 nodes apart, while message-passing networks fail. Importantly, when the subset of basis that accounts for global pooling is eliminated, the performance of Transformers drops similar to message-passing networks, showing their importance in modeling long-range interaction. Yet, simply having global basis is not enough, as seen in the failure of MLPπ. This indicates fine-grained interaction modeling via attention is essential even in this simple task.
Large-scale graph regression To further evaluate higher-order Transformers in large-scale setting, we used the PCQM4M-LSC dataset from Open Graph Benchmark [17], which is the largest graphlevel regression dataset composed of 3.8M molecular graphs. As test data is unavailable, we report the Mean Absolute Error (MAE) on validation dataset. In addition to the baselines from the benchmark, we also report the performances of second-order invariant MLP and a vanilla (first-order) Transformer4 for comparison. Further details can be found in Appendix A.2.5.
4As vanilla Transformer operates on node features only, we used Laplacian graph embeddings [2, 8] as positional embeddings so that the model can consider edge structure information.
The results are in Table 2. Second-order Transformer outperforms the message-passing GNNs (GCN, GIN) by a large margin, including the ones with a virtual node that can model long-range interactions (GCN-VN, GIN-VN). It suggests that higher-order attention is potentially better in handling longrange interactions on graphs than the current practice of augmenting GNNs with a virtual node. Furthermore, second-order Transformer outperforms invariant MLP, indicating that replacing sum-pooling with attention is important for scale-up. Finally, second-order Transformer significantly outperforms vanilla Transformer with Laplacian graph embeddings. This is presumably because node embeddings are insufficient to utilize features associated with edges, while second-order Transformers can naturally use all edge information. We also note that invariant MLP and vanilla Transformer have worse complexity than the second-order Transformer (O(m2) and O(n2), respec-
tively), while the second-order Transformer has O(m) complexity identical to GCN. Set-to-graph prediction An important advantage of our framework distinguished from most existing GNNs is that, by design, it can be applied to any learning scenario with different input and output orders (mixed-order). To demonstrate this, we tested higher-order Transformers in set-to-graph prediction tasks where the goal is to predict edge structure of a graph given a set of node features. We used two datasets following the prior work [29]. The first dataset Jets originates from particle physics experiments, where collision of high-energy particles gives a set of observed particles. The task is to partition the feature set of observed particles according to their origin. By viewing each subset of particles as a fully-connected graph, the problem is cast as a set-to-graph prediction. For the second dataset Delaunay, the task is to predict Delaunay triangulation [7] given a set of points in 2D space. Two datasets are used for this task, one containing 50 points and the other containing varying number of points ∈ {20, ..., 80}. For the baselines, we take the scores reported in Serviansky et. al. (2020) [29], which includes GNNs and invariant MLPs (S2G, S2G+) with L1→1 and L1→2. Our model is made by substituting the linear layers with Enc1→1 and Enc1→2 (mixed-order) respectively. Further details including the datasets, metrics, and baselines can be found in Appendix A.2.6. The results are outlined in Table 3. Mixed-order Transformers, both softmax and kernel attention, have favorable scores over all baselines. Especially, they outperform all baselines by a large margin in Delaunay: note that GNNs fall into trivial solution with high recall but very low precision. We particularly note that the Transformers’ performance in Delaunay (20-80) is comparable to Delaunay (50), with 0.3-0.7% drop in F1 score. Compared with S2G that exhibits ∼ 12% drop, this indicates attention mechanism within Enc1→2 is helpful in modeling varying number of nodes.
k-uniform hyperedge prediction One major advantage of our framework is that it naturally extends to higher-order data (hypergraphs). To demonstrate this, we consider higher-order extension of set-to-graph prediction task, where the goal is predicting k-uniform hyperedges (e.g., user-locationactivity) from node features. For evaluation, we used three datasets for transductive 3-edge prediction following the prior work [36]. The first dataset GPS derives from a GPS network [38], and contains (user-location-activity) hyperedges. The second dataset MovieLens is a social network dataset of tagging activities [13], containing (user-movie-tag) hyperedges. The third dataset Drug comes from a medicine network from FAERS5, containing (user-drug-reaction) hyperedges. As baselines, we consider higher-order invariant MLP (S2G+) and the state-of-the-art, self-attention based HyperSAGNN [36]. We implemented higher-order Transformer and S2G+ by substituting Enc1→2 and L1→2 in set-to-graph architectures to Enc1→3 and L1→3, respectively. Further details including the datasets, metrics, and baselines can be found in Appendix A.2.2 and Appendix A.2.7.
The results are in Table 4. Higher-order Transformer outperforms S2G+ in all datasets, and HyperSAGNN in all but one metric in MovieLens. The results suggest that higher-order self-attention is effective in learning higher-order representation beyond second-order graphs. The results are encouraging especially because we did not introduce any form of task-specific heuristics into the model, while some of the baselines such as Hyper-SAGNN depend on many inductive biases (static/dynamic branches, Hadamard power, etc.).

heading: 6 Discussion
text: In this paper, we proposed a generalization of Transformers to higher-orders, and applied a number of design strategies that reduce their complexity to a feasible level. Higher-order Transformers are attractive, both in theory and application. In theoretical aspect, it inherits the theoretical completeness and expressive power of invariant MLPs. In application aspect, it is potentially more powerful than message-passing neural networks due to global interaction modeling, and can be extended to a variety of useful mixed-order tasks involving sets, graphs, and hypergraphs.
At the same time, our work has some limitations that need to be addressed in future work. First, although complexity to input size can be lowered to linear, the number of basis grows rapidly with increasing order due to O((0.792k/ ln(k + 1))k) asymptotic formula of k-th Bell number [3], still making the model infeasible in higher orders. Improvement approaches such as finding a compact subset of basis that retains universality [29] and exploiting unorderedness [26] are promising in this direction. Second, our work builds upon tensor-based representation of graphs, which makes it difficult to be directly extended to hypergraphs containing edges with varying orders (e.g., co-citation networks). This is a common challenge to all tensor-based methods [14, 26, 25, 19, 29], and we believe addressing this would be an important future research direction.
Acknowledgements This work was supported in part by National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (2021R1C1C1012540, 2021R1A4A3032834) and Institute of Information & Communications Technology Planning & Evaluation (IITP) grant (2021-0-00537, 2019-0-00075).
5https://www.fda.gov/Drugs/

Review summaries:
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. 

The authors propose a framework to generalize transformers to any order permutation invariant data, especially hypergraphs. While the generalization itself is incremental compare to [20], the naïve approach gives unbearable computational complexity. Hence, the authors propose a series of tricks to lower the computational complexity which is their main contribution.