摘要

我们提出了一种将Transformer模型推广到任意阶排列不变数据（集合、图和超图）的方法。首先，我们观察到Transformer可以视为DeepSets的泛化，即一阶（集合输入）排列不变多层感知机（MLP）。接着，基于最近研究中描述的高阶不变MLP，我们将自注意力机制的概念扩展到更高阶，并提出了适用于$k$阶数据的高阶Transformer模型（对于图数据，$k=2$；对于超图数据，$k>2$）。然而，高阶Transformer的复杂度为$\mathcal{O}(n^{2k})$，随着输入节点数$n$的增加而变得难以承受。为了解决这一问题，我们提出了一种稀疏高阶Transformer模型，其复杂度与输入超边数呈二次关系，并进一步采用了核注意力方法将其复杂度降低至线性。特别是，我们证明了采用核注意力的稀疏二阶Transformer在理论上比消息传递操作更具表达能力，同时具有渐近相同的复杂度。我们的模型在大规模图回归和集合到（超）图预测任务中显著优于不变MLP和消息传递图神经网络。我们的实现代码已发布在https://github.com/jw9730/hot。