Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs

TL;DR

This paper extends Transformers to higher-order permutation-invariant data like graphs and hypergraphs, introduces sparse and kernel attention methods for scalability, and demonstrates superior performance over existing models.

Contribution

We propose higher-order Transformers for complex data structures, reducing computational complexity and enhancing expressiveness compared to prior invariant models.

Findings

Sparse higher-order Transformers are more expressive than message passing GNNs.

Kernel attention reduces complexity to linear in input size.

Models outperform invariant MLPs and GNNs in large-scale graph tasks.

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 complexity $\mathcal{O}(n^{2k})$ 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.