NeurIPS 2025｜从层次化掩码的视角统一并增强 Graph Transformer

作者：邢宇杰，王啸，伍斌，黄海，石川单位：北京邮电大学；北京航空航天大学代码： https://github.com/null-xyj/M3Dphormer

引言

作为一种基础的数据结构，图被广泛用于表示现实世界系统中复杂而多样的交互关系，例如社交网络和脑网络。为了学习高质量的节点表征，人们提出了多种图神经网络（GNN）。然而，由于消息传递机制所固有的强局部性归纳偏置，这类方法的性能受到了限制。受到变换器模型在多个机器学习领域中成功应用的启发，将变换器架构引入图结构成为一种有前景的方向，因为其具备在更大范围内建模交互的能力。 Graph Transformer 利用变换器架构的核心组件——多头注意力机制，自适应地建模多样的节点交互并学习具有表达力的表征。一类重要研究将整个图视为完全连接，并通过注意力机制捕获节点之间的成对依赖关系。另一类方法通常通过节点采样或特征聚合为每个节点构建一个标记序列，再使用变换器来捕获多尺度交互。此外，还有一些研究利用图划分来实现高效的交互建模并学习高质量表征。尽管已有多种 Graph Transformer 被提出，它们往往依赖于针对特定交互类型而精心设计的模型结构，从而限制了对其他重要交互的建模能力。这引出了一个自然的问题：**是否存在一种统一的图变换器视角，能够灵活地对多样的节点交互进行建模？**为了回答这个问题，我们从图中不同层级交互信息建模的角度出发，提出了一个基于层次化掩码的 Graph Transformer 统一框架。基于此框架，我们对理想态的图注意力掩码进行了理论分析，发现其核心设计原则为：一个有效的图注意力掩码应当确保充分大的感受野以及较高的标签一致率。我们还进一步提出 MDphormer 方法，其设计了三种有效的图注意力掩码来捕获不同层次的关联信息，采用两级注意力专家路由机制实现不同层次关联信息的自适应整合并提出双重注意力计算模式大幅提升模型的计算效率。

预备知识

我们将一个带属性的图记作，其中表示包含个节点的节点集合，表示包含条边的边集合，为节点特征矩阵，其中表示节点的维特征向量。邻接矩阵记为，其中若节点与节点之间存在一条边，则，否则。在节点分类任务中，标签集合记为，节点标签表示为，其中为节点的 one-hot 标签。整个节点集合可以划分为训练集、验证集和测试集。 Graph Transformer：

Graph Transformer 的核心组件是多头注意力机制（MHA），其形式如下：其中，表示输入节点表征，为节点数量，为表征维度。是第个注意力头的可训练投影矩阵，为每个头的隐藏维度。注意力分数矩阵由掩码矩阵进行筛选，其中表示节点可以对节点施加注意力。在掩码操作中，所有无效位置的注意力分数会被置为，从而在 Softmax 之后其贡献变为零。所有注意力头的输出会被拼接形成最终输出：

通过统一的层次化掩码框架重新审视 Graph Transformers

图中的交互通常具有层次化结构，包括局部连接、簇级关系以及全局关联。每一层级都为有效的图表征学习提供了关键信息。图 1 展示了层次化交互的重要性示例，其中节点标签用不同颜色表示。由于局部同质性，节点可以被准确分类，因为其大多数邻居具有相同的标签。然而，仅依靠这种局部交互不足以分类节点和。通过利用簇级交互，节点可以被正确识别，因为它位于一个结构连贯的簇（蓝色区域）中。此外，若假设绿色标签节点的表征分布与蓝色和黄色节点显著不同，则可以自适应地学习全局交互（由虚线表示），进一步提升节点的分类效果。

基于上述动机，我们提出了一个统一的层次化掩码框架，使 Graph Transformers (GTs) 能够一致地建模多层级的交互关系。

Graph Transformers 的统一层次化掩码框架

为了形式化这一框架，我们将节点交互划分为三类：（1）N–N：节点与节点之间的交互；（2）N–S：节点与节点集合之间的交互；（3）S–S：节点集合之间的交互。对于 N–N，我们直接设置表示节点到节点的连接。对于 N–S 和 S–S，我们将每个节点集合视为一个虚拟超节点，并将节点集合扩展为。这使得 N–S 与 S–S 均可以等价地建模为 N–N，从而在不同交互层级之间实现统一而灵活的掩码设计。在这一统一框架的基础上，我们进一步说明现有的 GT 是如何在局部、簇级和全局层面上通过其注意力掩码隐式建模层次化交互的。 局部交互：

建模局部交互通常涉及捕获来自跳邻域的信息。一种方法是显式建模目标节点与其跳邻居之间的 N–N 交互，对应于掩码。混合架构（GNN–Transformer）通过 GNN 模块隐式采用掩码，并在层间递归聚合跳信息。部分 token 化的 GT 则将多跳节点特征聚合作为 Transformer 的输入 token，从而等价地应用掩码集合。 簇级交互：

一些近期方法通过将图划分为若干不相交簇来捕获簇级交互。我们定义划分函数表示节点被分配到簇，反向划分函数表示属于簇的节点集合。将每个簇视为一个虚拟超节点后，我们将这些节点集合记为。一种方法侧重于建模簇与簇之间的交互（S–S 交互），对应于掩码，其中当时，。另一类方法建模更细粒度的 N–S 交互，使得每个簇能够关注所有真实节点，对应掩码，其中当且时，。为了区分不同簇的贡献，注意力分数会根据簇之间的连通性进行调制，并通过由隐式引出的簇注意力进一步细化。也有方法专注于簇内的 N–N 交互，即在每个簇内部隐式应用掩码，其中当时，，并额外通过捕获簇间交互。 全局交互：

一种常见策略是将整个图视为完全连接，从而捕获全局 N–N 交互，对应于全局掩码。另一类方法通过引入一组虚拟超节点并使其与所有真实节点双向连接，从而通过 N–S 方式近似全局依赖性。这对应于掩码，其中当且，或且时，。我们在表 5 中汇总了已有 Graph Transformer 方法中对应的图注意力掩码。

理论分析

为了研究层次化掩码在节点分类中的不同贡献，我们基于类别条件表征模型构建了一个理论框架。具体而言，令为一个具有标签集合的图，并假设节点标签在整体上均匀分布。对于标签为的节点，其初始表征从一个维高斯分布中采样：其中满足，且各类的原型向量两两正交，即对于所有，均有。令节点的真实标签为，其感受野由掩码向量指定，其中包含个非零元素。定义为该感受野中标签为的节点所占比例，为这些节点获得的平均注意力权重。记为节点经注意力更新后的表征。我们考虑一个基于相似度的分类器，当满足时预测为正确的标签，其中为模型隐式学习到的决策阈值。

定理3.1 表明，正确分类的概率主要取决于以下因素：、、以及方差集合。由于和分别受到训练动态与输入分布的影响，我们将分析重点放在由注意力掩码决定的与上。其显示，当和的取值更大时，正确分类概率的上下界均更高。由此可得出一个构造掩码的基本原则：**有效的注意力掩码应确保足够大的感受野以及较高的标签一致性。**基于上述定理，我们进一步在若干典型情境下分析现有 GT 所采用的层次化掩码的适用性： 1. **强局部同质性的节点：**当节点具有显著的局部同质性时，使用局部掩码（如、）是有效的，因为它们通常带来较大的。 1. **位于簇边界的节点：**在簇边界附近，局部同质性减弱，因为部分邻居属于其他簇。在这种情况下，簇级掩码（如）能够提供更高的和更大的，从而可能提升分类性能（前提是图划分较为准确）。 1. **在簇内属于少数类的异质性节点：**对于这类节点，当采用局部掩码或簇级掩码时，往往会非常小，甚至低于。此时，全局掩码更为合适，因为它保证，并提供的大感受野。相比之下，的效果可能较弱，因为全局虚拟节点缺乏明确的语义标签。 1. **通过簇间注意力近似全局交互的情形：**有些研究通过簇间交互（如）近似全局依赖。然而根据定理3.1，簇级表征会被簇内多数类主导。在随后的簇间聚合中，这种偏置会进一步放大，从而使少数类节点更容易被误分类为多数类。 1. **类表征分布明确（即较小）的节点：**对于这类节点，注意力权重可以被有效学习为的近似形式，使得无论使用何种掩码，其正确分类概率都较高。

综上所述，没有单一的掩码能够在所有场景下持续满足上述原则。然而，不同层级的层次化掩码在节点分类中具有互补优势，将它们结合起来可自然地符合这一基本原则。

实验分析

为了进一步研究在真实场景中结合多层级交互掩码是否能够带来收益，我们构建了三个分别使用局部掩码、簇级掩码和全局掩码的 GT，并在七个节点分类数据集上独立训练。随后，我们采用三种集成策略对它们的输出进行融合：Mean：对预测概率取平均；Max：选择最大预测概率；Oracle：一种理想化的上界，始终选择最优预测结果。单模型与集成方法的性能如表1所示：

根据表1可得到三点关键观察：1）Oracle 的强表现表明，正确整合来自不同交互层级的互补信息能够带来显著的性能提升。2）简单的集成方法（Mean 与 Max）在 7 个数据集中有 5 个表现不如最佳的单一掩码模型，这揭示了 Graph Transformer 面临的一个核心挑战：如何有效整合多层级交互信息？此外，我们的实验还揭示了一个重要的效率问题：3）即使是仅包含 2 层、2 个头且只有单一掩码的 Transformer，在 PubMed 上也会消耗 21 GB 的 GPU 显存。尽管已有多种高效 Transformer 变体（如基于核函数的线性注意力方法和 FlashAttention）能够缓解的复杂度，但由于图掩码模式的不规则性与多样性，这些方法在图领域的适用性依然有限。因此，引出了另一个核心挑战：如何在大规模图上高效实现具备不规则掩码的 GT？

方法：MDphormer

我们提出了 MDphormer，这是一种基于混合专家的 Graph Transformer，具备多级掩码机制与双重注意力计算模式。其整体框架如图2所示。在理论分析的指导下，我们首先设计了三个层次化注意力掩码，用于全面建模多层级的交互（见部分 c）。为了自适应地整合不同交互层级的信息，我们引入了双层注意力专家路由机制，其中每个专家对应一个绑定特定掩码的 MHA 模块（见部分 b）。此外，我们设计了双模注意力计算方案，以确保计算效率（见部分 d）。

整体框架

我们首先介绍 MDphormer 的整体架构。初始表征通过如下公式计算：其中为可学习的线性投影矩阵。随后，模型堆叠层 MDphormer 层，第层的计算定义如下：其中，为上一层的输入，为残差投影矩阵，为归一化函数，为激活函数。双层注意力专家路由机制基于层次化掩码集合来整合不同交互层级的信息。在经过层之后，使用线性分类器获得最终预测：其中为分类器权重矩阵。模型通过交叉熵损失进行优化，计算范围包括训练节点以及下一小节将介绍的标签特定全局虚拟节点。

理论指导下的层次化掩码设计策略

为了实现对多层级节点交互的全面建模，我们基于理论分析提出了一种层次化掩码集合的设计策略。 **Local mask design（局部掩码设计）：**我们采用作为局部掩码，原因如下：

如相关研究所示，随着跳数的增加，同质性比例会快速下降。根据 Theorem~，较低的会导致更低的分类概率。
忽略了目标节点与其跳邻居之间的距离信息，因此需要额外的基于距离的编码；相较之下，可通过跨层的递归聚合隐式捕获距离信息。
相比更加稀疏，使后续引入的双模注意力计算方案能够更高效地执行。

**Cluster mask design（簇级掩码设计）：**我们首先将图划分为个不相交簇，并引入一组簇级虚拟节点。随后定义新的簇级掩码，其规则为：当满足以下任一条件时，令： (i)且；(ii)且。其中，与分别表示此前章节中定义的簇划分函数与反向划分函数。该设计使注意力仅限于同一簇内的节点–簇级交互。每个虚拟节点的特征通过对其对应簇中所有真实节点特征取平均获得。与相比，所提出的将非零元素比例从显著降低至，由于通常，计算效率大幅提升。此外，如命题4.1所示，所捕获的交互能够被有效近似。

尽管这需要额外的一层，但由于稀疏度的大幅降低，我们提出的双重注意力计算模式仍能够带来显著的计算效率提升。 **Global mask design（全局掩码设计）：**我们提出了新的全局掩码，它在的基础上显式融入了标签语义。具体而言，我们新增个全局虚拟节点，索引为其中每个虚拟节点对应一个类别标签。掩码在以下任一条件满足时取值为 1： (i) 且； (ii) 且在这种结构下，每个真实节点都可以关注所有全局节点，而每个全局节点仅从训练集中对应标签的节点聚合信息。与类似，中节点的特征通过对其对应类的训练节点特征进行平均得到。设为类别的全局虚拟节点，为训练集中属于类别的节点数。根据定理3.1，其更新后的表征满足：因为，对所有有，且注意力权重。方差降低至表明其分布更加集中于类别均值附近。根据定理3.1，这将带来更优的正确分类概率界。

两级注意力专家路由机制

为了自适应地整合来自不同交互层级的信息，我们提出了双层注意力专家路由机制，这是 MDphormer 的核心组件。每个专家对应一个绑定特定注意力掩码的 MHA 模块。基于表1中的经验观察，即局部掩码在大多数场景下表现最佳，我们在第一级路由中优先选择局部专家；第二级则在簇级专家与全局专家之间进一步细化选择。双层路由机制形式化定义如下：其中，表示层次化掩码集合，为我们提出的双模注意力计算方案。是用于第一级和第二级路由的可学习门控参数。Sigmoid 函数将门控值和限制在区间内。为了突出局部交互在经验上的重要性，我们将和均初始化为零向量，使每个节点在训练初期的路由权重为，从而优先关注局部注意力。基于表1的进一步观察，即所有交互层级均对分类任务具有重要贡献，我们不采用 top-选择，而是使用学习得到的路由权重、和对所有专家的输出进行加权聚合。

双重注意力计算模式

最后，我们提出双模注意力计算方案以提升计算效率。由于图掩码的非规则性，现有高效注意力变体难以直接适用；但图掩码本身的稀疏性为另一条优化路径提供了可能——稀疏注意力计算。与标准的稠密注意力不同，后者需在施加二值掩码前构造完整的注意力矩阵，而稀疏注意力仅对满足的有效节点对计算注意力得分。对稀疏注意力实现进行分析可知，其空间复杂度为，其中为掩码的非零元素数量，为注意力头数，为每个头的隐藏维度。尽管该方法已将复杂度从显著降低至，但常数项在实际中仍不可忽略。为进一步提升效率，我们提出双模注意力计算方案，根据注意力掩码的局部稀疏性在稠密计算与稀疏计算之间动态切换。具体而言，我们将注意力掩码划分为个互不重叠的区域，其中每个区域由查询节点集合及其对应的键节点集合定义。对于每个区域，我们根据命题4.2选择最优的计算模式。

该结果为计算模式的动态选择提供了理论指导：在局部稀疏度较高时使用稀疏计算，而在区域变得足够稠密时切换为稠密计算。基于此，我们将双模注意力计算方案形式化为：

实验

节点分类

表2展示了节点分类任务的结果。主要观察如下：

MDphormer 在 9 个数据集上均优于所有基线模型，体现了其更强的交互建模能力。
与传统 GNN 和基于 MoE 的 GNN 相比，MDphormer 在捕获层次化交互方面展现出显著优势。
与 GT 类基线相比，MDphormer 取得了明显提升，这验证了两点：其一，全面建模多层级交互是有效的；其二，双层注意力专家路由机制能够有效自适应整合多层级信息。

消融实验

我们进行了消融实验，以从有效性和效率两个方面评估 MDphormer 的表现。首先，我们构建了五个消融变体，分别通过以下方式实现：1）移除单个专家模块；2）禁用注意力专家路由机制；3）将双层注意力路由机制替换为单层门控方案。

表3的结果表明：1）移除任意一个专家都会导致性能下降，这强调了全面建模层次化交互的重要性；2）禁用路由机制会带来显著退化，说明简单的聚合操作不足以有效整合多层级交互信息；3）单层路由变体与 MDphormer 之间的性能差距表明，所提出的双层注意力专家路由机制具有明显优势。

随后，我们在图3 中报告了 MDphormer 及其采用稠密和稀疏计算方案的两个变体的 GPU 显存占用情况。结果显示，稠密方案占用的显存最高，并在四个数据集上出现了内存溢出（OOM）错误；稀疏方案虽然显著降低了显存消耗，但在 Ogbn-Arxiv 上仍无法运行。相比之下，我们提出的双重注意力计算方案在显存效率方面表现更优，并能够成功扩展到所有评测图数据集上。

参数分析

MDphormer 中唯一的重要超参数是簇数量，它会影响簇级掩码的质量。对于每个数据集，我们根据其规模选择合适的范围。图4展示了模型在不同取值下的性能表现。总体来看，MDphormer 在大多数数据集上对的选择表现出较强的鲁棒性。唯一的例外是 Chameleon，一个仅包含 890 个节点的小型图，其划分质量会随着的变化而产生较大波动，从而导致更明显的性能差异。