Paper Reading：Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework

目录

• 研究动机
• 文章贡献
• 预备知识
• 本文方法
  • 统一层次化掩码框架
  • 理论分析
  • 层次化掩码设计策略
  • 双层注意力专家路由机制
  • 双注意力计算方案
  • 整体架构
• 实验结果
  • 数据集和实验设置
  • 节点分类性能对比
  • 消融实验
  • 内存效率分析
  • 参数敏感性分析
• 优点和讨论

Paper Reading 是从个人角度进行的一些总结分享，受到个人关注点的侧重和实力所限，可能有理解不到位的地方。具体的细节还需要以原文的内容为准，博客中的图表若未另外说明则均来自原文。

论文概况	详细
标题	《Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework》
作者	Yujie Xing, Xiao Wang, Bin Wu, Hai Huang, Chuan Shi
发表会议/期刊	NeurIPS (Conference on Neural Information Processing Systems)
发表年份	2025
会议/期刊等级	CCF-A
论文代码	https://github.com/null-xyj/M3Dphormer

作者单位：

1. Beijing University of Posts and Telecommunications, China（北京邮电大学）
2. Beihang University, China（北京航空航天大学）

研究动机

图神经网络基于消息传递机制，具有局部性的归纳偏置，在建模长距离交互时存在固有局限。受 Transformer 在多个机器学习领域成功的启发，图Transformer（GT）通过多头注意力机制自适应建模多样化的节点交互，成为图表示学习的一个重要方向。现有的图Transformer方法种类繁多，但各自依赖针对特定交互类型设计的复杂架构，缺乏统一的视角：

• 一类方法将整个图视为全连接，直接应用注意力机制捕获节点对依赖关系。
• 另一类方法为每个节点构建 token 序列，通过节点采样或特征聚合后使用 Transformer 捕获多尺度交互。
• 还有一类方法利用图划分实现高效的交互建模。

这些方法分别侧重局部、簇或全局层面的交互，架构设计差异很大，难以灵活适应不同的交互模式。图中的交互天然呈现层次化组织，包括局部连接、簇关系和全局关联。每个层次的信息对有效的图表示学习都不可或缺。下图给出了一个例子，局部同配性可以正确分类节点 u1，但不足以分类 u2 和 u3。簇级交互可以正确识别 u2，全局交互可以进一步增强 u3 的分类。这说明不同层次的交互信息是互补的，单一层次的建模无法覆盖所有情况。

从理论角度看，现有工作缺乏对不同层次掩码如何影响节点分类性能的系统分析。不同掩码的适用场景、设计的基本原则、以及如何有效整合多层次掩码等问题都没有得到充分解答。从实践角度看，简单的集成方法（如均值和最大池化）往往不如最优的单掩码模型，如何有效融合多层次交互信息是一个挑战。此外，图Transformer在中等规模图上的内存消耗也构成了重要的效率瓶颈。

综上所述，现有图 Transformer 方法缺乏统一的理论框架，多层次交互的有效整合和计算效率是两个亟待解决的问题。本文针对这些不足，提出了统一的层次化掩码框架，并基于该框架设计了 M3Dphormer 模型。

文章贡献

针对现有图 Transformer 架构分散、缺乏统一理论视角的问题，本文提出了一个统一的层次化掩码框架，揭示了模型架构与注意力掩码构建之间的内在等价关系。该框架将节点交互统一为节点-节点（N-N）、节点-集合（N-S）和集合-集合（S-S）三种类型，通过引入虚拟超节点将后两种等价转化为 N-N 交互，从而可以通过精心设计的注意力掩码一致地建模不同层次的交互。在该框架下，本文进行了理论分析，证明正确分类概率的上下界与感受野大小和标签一致性程度正相关，由此得出注意力掩码的基本设计原则：有效的注意力掩码应同时保证足够大的感受野和较高的标签一致性。分析表明，单一掩码无法在所有场景下满足这一原则，不同层次的掩码具有互补优势。在此基础上，本文提出 M3Dphormer，一种基于混合专家的图 Transformer，具有多层次掩码和双注意力计算。M3Dphormer 采用三个层次化掩码（局部、簇、全局），通过双层专家路由机制自适应整合多层次交互信息，并引入双注意力计算方案根据局部掩码稀疏性动态切换稠密和稀疏模式以保证可扩展性。在九个基准数据集上的广泛实验表明，M3Dphormer 取得了 SOTA 性能，验证了统一框架和模型设计的有效性。

预备知识

理解本文方法需要了解两个概念。图 Transformer 与注意力掩码。 图 Transformer 的核心组件是多头注意力（MHA），给定输入表示 \(H \in \mathbb{R}^{N \times d}\)，第 \(i\) 个注意力头的计算为：

\[\text{head}_i(H, M) = \text{Softmax}\left(\text{Mask}\left(\hat{A}^{(i)}, M\right)\right) V^{(i)} \]

其中 \(\hat{A}^{(i)} = \frac{Q^{(i)} K^{(i)\top}}{\sqrt{d_h}}\) 是注意力得分矩阵，\(Q^{(i)} = H W_Q^{(i)}\)、\(K^{(i)} = H W_K^{(i)}\)、\(V^{(i)} = H W_V^{(i)}\) 分别是查询、键和值的投影。\(M \in \{0, 1\}^{N \times N}\) 是注意力掩码，\(M_{u,v} = 1\) 表示从节点 \(u\) 到 \(v\) 的注意力有效，\(M_{u,v} = 0\) 时对应注意力得分被设为 \(-\infty\)，经过 Softmax 后贡献为零。所有头的输出拼接得到最终结果。

图的层次化交互。 图中的交互可以分为三个层次，三个层次各有适用场景，没有单一层次在所有情况下都最优。

交互层次	定义	实现方式
局部交互	节点与其 \(K\) 跳邻居之间的连接	局部同配性假设
簇级交互	同一簇内节点之间的关系	图划分算法（如 METIS）
全局交互	整个图范围内节点之间的关联	全连接注意力或虚拟节点

本文方法

M3Dphormer 的整体框架如图所示，包含三个组件：理论指导的层次化掩码设计、双层注意力专家路由机制和双注意力计算方案。

统一层次化掩码框架

本文首先提出一个统一框架来理解和设计图Transformer。节点交互分为三类：

交互类型	描述
N-N	节点与节点之间的交互，直接设置 \(M_{u,v} = 1\)
N-S	节点与节点集合之间的交互
S-S	节点集合与节点集合之间的交互

对于 N-S 和 S-S 交互，将每个节点集合视为一个虚拟超节点 \(v'\)，扩展节点集合为 \(V = V \cup \{v'\}\)，这样 N-S 和 S-S 都可以等价建模为 N-N 交互。基于这一统一框架，现有图 Transformer 可以被重新解释为对应特定掩码的设计：

• 局部交互掩码。 GOAT 通过目标节点与其 \(K\) 跳邻居之间的 N-N 交互建模，对应掩码 \(M_{l1} = A^K\)。GNN-Transformer 混合架构通过 GNN 模块隐式采用掩码 \(M_{l2} = A\)，在层间递归聚合 \(K\) 跳信息。一些 token 化的 GT 将多跳节点特征聚合为 Transformer 的输入 token，实际上应用了一组掩码 \(\{A^k : 1 \leq k \leq K\}\)。

• 簇级交互掩码。 基于图划分的 GT 将图划分为不相交的簇。Graph ViT 建模簇之间的 S-S 交互，对应掩码 \(M_{c1}\)。Cluster-GT 建模更细粒度的 N-S 交互，每个簇关注所有真实节点，对应掩码 \(M_{c2}\)。CoBFormer 关注每个簇内的 N-N 交互，对应掩码 \(M_{c3}\)，并额外应用 \(M_{c1}\) 捕获簇间交互。

• 全局交互掩码。 全连接注意力对应全局掩码 \(M_{g1} = \mathbf{1}_{N \times N}\)。Exphormer 通过引入一组与所有真实节点双向连接的全局虚拟超节点来近似全局依赖，对应掩码 \(M_{g2}\)。

理论分析

本文基于类别条件表示模型建立了理论分析框架。假设节点初始表示服从高斯分布 \(z \sim \mathcal{N}(\mu_c, \sigma_c^2 I)\)，类别原型正交。设节点 \(u\) 的真实标签为 \(c\)，其感受野由掩码向量 \(M_{u,:}\) 指定，包含 \(k\) 个非零条目。定义 \(\rho_{c'}\) 为感受野内标签为 \(c'\) 的节点比例，\(\alpha_{c'}\) 为分配给这些节点的平均注意力权重。定理3.1 给出了节点更新表示的分布和正确分类概率的上下界：

\[\hat{z}_u \sim \mathcal{N}\left(\sum_{i=1}^{|C|} k\rho_i\alpha_i\mu_i, \sum_{i=1}^{|C|} k\rho_i\alpha_i^2\sigma_i^2 I\right) \]

正确分类概率的上下界均关于 \(k\)（感受野大小）、\(\rho_c\)（标签一致性）和 \(\alpha_c\)（注意力权重）单调递增，关于各类别方差递减。定理说明：有效的注意力掩码应同时保证足够大的感受野和较高的标签一致性。基于此定理，本文进一步分析了不同层次掩码在典型场景下的适用性：

1. 对于局部同配性强的节点，局部掩码有效，因为 \(\rho_c\) 通常较大
2. 对于簇边界附近的节点，局部同配性减弱，簇掩码可以提供更高的 \(\rho_c\) 和更大的 \(k\)
3. 对于簇内少数标签的异配节点，局部或簇掩码的 \(\rho_c\) 可能很低，全局掩码 \(M_{g1}\) 更优
4. 通过簇间注意力近似全局交互可能放大多数类偏差，增加少数类节点被误分类的风险
5. 对于表示分布明确的类别（方差小），注意力权重可以有效学习，掩码选择的影响较小

分析表明，没有单一掩码能在所有场景下一致满足设计原则。不同层次的掩码在节点分类中具有互补优势，整合多层次掩码是遵循设计原则的自然途径。

层次化掩码设计策略

受理论分析指导，本文设计了三个层次化注意力掩码 \(M = \{M_{l2}, M_{c4}, M_{g3}\}\)。

• 局部掩码。 采用 \(M_{l2} = A\) 作为局部掩码。相比 \(M_{l1} = A^K\)，\(M_{l2}\) 有三个优势：同配率 \(\rho_c\) 随跳数 \(K\) 增加快速下降，低阶跳数的 \(\rho_c\) 更高；\(M_{l2}\) 可以通过层间递归聚合隐式捕获距离信息，无需显式位置编码；\(M_{l2}\) 更稀疏，计算效率更高。

• 簇掩码。 首先使用 METIS 将图划分为 \(P\) 个不相交的簇，引入一簇级虚拟节点集合 \(V_p\)。定义新的簇掩码 \(M_{c4}\)，其中 \(M_{c4_{u,v}} = 1\) 当且仅当：(i) \(u \in V\) 且 \(v \in \{u, N+P(u)\}\)，或 (ii) \(u \in V_p\) 且 \(v \in P^{-1}(u-N)\)。相比 \(M_{c3}\)，\(M_{c4}\) 的非零率从 \(1/P\) 大幅降低到 \(3N/(N+P)^2\)（因 \(P \ll N\)），计算效率更高。命题4.1 证明，使用 \(M_{c3}\) 的单层 GT 可以等价地由两层使用 \(M_{c4}\) 的 GT 建模，尽管增加了一层，但稀疏度降低带来的计算节省仍然显著。

• 全局掩码。 提出新的全局掩码 \(M_{g3}\)，在 \(M_{g2}\) 基础上显式融入标签语义。添加 \(|Y|\) 个全局虚拟节点，每个对应一个类别标签。\(M_{g3_{u,v}} = 1\) 当且仅当：(i) \(u \in V\) 且 \(v \in V_g\)，或 (ii) \(u \in V_g\) 且 \(v \in \{t \in V_{\text{train}} : y_t = y_u\}\)。每个真实节点可以关注所有全局节点，每个全局节点只从对应标签的训练节点聚合信息。根据定理3.1，全局虚拟节点的表示集中在类别均值附近（方差降低），从而改善正确分类概率的上下界。

双层注意力专家路由机制

为自适应整合不同交互层次的信息，本文提出双层注意力专家路由机制作为 M3Dphormer 的核心组件。每个专家对应一个配备特定注意力掩码的 MHA 模块。受实验观察启发，局部掩码在大多数情况下表现最优，因此第一级路由优先选择局部专家。第二级在簇专家和全局专家之间进行细化选择。双层路由机制的形式化定义为：

\[\text{BiMoE}(H, M) = g_1 \cdot \text{MHA}_D(H, M_{l2}) + g_2 \cdot \text{MHA}_D(H, M_{c4}) + g_3 \cdot \text{MHA}_D(H, M_{g3}) \]

其中门控权重通过两层 sigmoid 门控计算：

\[g_1 = \beta_1, \quad g_2 = (1-\beta_1) \cdot \beta_2, \quad g_3 = (1-\beta_1) \cdot (1-\beta_2) \]

\[\beta_1 = \text{Sigmoid}(H W_G^1), \quad \beta_2 = \text{Sigmoid}(H W_G^2) \]

\(W_G^1, W_G^2 \in \mathbb{R}^{d \times 1}\) 是可学习门控参数。为强调局部交互的重要性，两者初始化为零向量，初始路由权重为 \([0.5, 0.25, 0.25]\)，在训练早期优先考虑局部注意力。受所有交互层次都对分类有显著贡献的观察启发，模型使用学习的路由权重聚合所有专家的输出，不采用 top-k 选择。

双注意力计算方案

图掩码的不规则性阻碍了高效注意力变体的应用，但其固有的稀疏性开辟了稀疏注意力计算的优化路径。与标准稠密注意力先构建完整注意力矩阵再应用二值掩码不同，稀疏注意力仅为掩码中有效的节点对计算注意力得分。稀疏注意力的空间复杂度为 \(O(6mHd_h)\)，其中 \(m\) 是掩码中非零条目的数量。

命题4.2 给出了稀疏方案与稠密方案的效率比较条件：当区域 \(R_i\) 内的非零率 \(\kappa_{R_i} < \frac{1}{3d_h}\) 时，稀疏注意力方案比稠密方案更高效。基于理论指导，双注意力计算方案将图划分为多个区域，每个区域根据局部稀疏度动态选择稀疏或稠密计算模式，最后聚合所有区域的注意力输出：

\[\text{MHA}_D(H, M) = \text{Comb}_{i=1}^{K}\left(\text{SelectMode}(R_i)(H, R_i)\right) \]

其中 \(\text{SelectMode}(R_i)\) 表示对区域 \(R_i\) 选择稀疏或稠密计算模式，\(\text{Comb}_{i=1}^{K}(\cdot)\) 聚合所有 \(K\) 个分区的注意力输出。

整体架构

M3Dphormer 的整体架构采用 Pre-RMSNorm 和 ReLU 激活函数。初始表示由线性投影给出 \(H_0 = X W_{\text{in}}\)。模型堆叠 \(L\) 个 M3Dphormer 层，第 \(l\) 层的计算为：

\[H_l = \text{ACT}\left(\text{BiMoE}_l(\text{Norm}_l(H_{l-1}), M)\right) + H_{l-1} W_{\text{res}}^l \]

其中 \(W_{\text{res}}^l\) 是残差投影矩阵，\(\text{Norm}_l(\cdot)\) 是归一化函数，\(\text{ACT}(\cdot)\) 是激活函数。\(L\) 层后应用线性分类器得到最终预测。模型使用交叉熵损失优化，损失计算覆盖训练节点和标签特定的全局虚拟节点。

实验结果

数据集和实验设置

本文在九个数据集上评估 M3Dphormer，数据集涵盖引文网络、社交网络和网页网络等多个领域：

数据集	图类型
Cora	同配性图
Citeseer	同配性图
Pubmed	同配性图
Computer	同配性图
Photo	同配性图
Ogbn-Arxiv	同配性图
Squirrel	异配性图
Chameleon	异配性图
Minesweeper	异配性图

对比方法分为五类，评价指标为节点分类准确率，Minesweeper 数据集使用 ROC-AUC。

类别	方法
经典 GNN	GCN、GAT、GraphSAGE
增强版经典 GNN	GCN、GAT、SAGE*
先进 GNN	GPRGNN、FAGCN
SOTA 图 Transformer	NAGphormer、Exphormer、SGFormer、CoBFormer、PolyNormer
基于 MoE 的 GNN	Mowst、GCN-MoE

节点分类性能对比

节点分类结果如表所示。M3Dphormer 在全部九个数据集上一致优于所有基线方法，展示了其交互建模能力。与传统 GNN 和基于 MoE 的 GNN 相比，M3Dphormer 通过全面捕获层次化交互展现出明显优势。以 Squirrel 和 Chameleon 这两个异配性较强的数据集为例，M3Dphormer 的准确率分别达到 44.34% 和 47.09%，比最强 GNN 基线（GCN-MoE 的 43.02% 和 44.57%）分别提升 1.32 和 2.52 个百分点。在 Minesweeper 数据集上，M3Dphormer 的 ROC-AUC 达到 98.27%，比 GCN* 的 97.39% 提升 0.88 个百分点。

与图 Transformer 基线相比，M3Dphormer 同样表现出显著提升。在异配性数据集上提升尤为明显，Squirrel 上比最强 GT 基线（Mowst 的 41.72%）提升 2.62 个百分点，Chameleon 上比最强基线（PolyNormer 的 44.30%）提升 2.79 个百分点。在同配性数据集上也有稳定提升，如 Computer 上达到 92.09%，比 CoBFormer 的 91.64% 提升 0.45 个百分点。在 Ogbn-Arxiv 大规模数据集上达到 73.54%，比 PolyNormer 的 73.27% 提升 0.27 个百分点。这些结果验证了全面交互建模的收益以及双层注意力专家路由机制在自适应整合多层次信息方面的有效性。

消融实验

本文进行了消融实验以验证各组件的作用，结果如表所示。消融设置包括：移除局部专家、移除簇专家、移除全局专家、移除路由机制（简单平均）、将双层路由替换为单层门控。实验得出三个主要结论：

1. 移除任何一个专家都会导致性能一致下降，证明全面建模层次化交互的必要性。其中移除局部专家造成的性能下降最大，在多个数据集上下降 3 到 6 个百分点，说明局部交互是最基础的信息来源。移除簇专家和全局专家也有明显影响，幅度相对较小。
2. 禁用路由机制（简单平均聚合）导致显著的性能下降，说明简单聚合不足以有效整合多层次交互信息。
3. 单层路由变体与 M3Dphormer 之间的性能差距证明了双层注意力专家路由机制的优势。双层设计通过先确定局部专家的权重再在簇和全局专家间分配，更符合交互层次的固有结构。

内存效率分析

本文比较了 M3Dphormer 及其稀疏和稠密计算变体的 GPU 内存使用情况，结果如图所示。稠密方案内存消耗最高，在四个数据集上出现内存溢出（OOM）。稀疏方案大幅降低了内存使用，但在 Ogbn-Arxiv 上仍无法运行。双注意力方案在所有评估的图上都成功运行，取得了最优的内存效率。

结果验证了双注意力计算方案的有效性。通过根据局部稀疏度动态选择计算模式，模型在稀疏区域享受稀疏计算的内存优势，在相对稠密的区域使用稠密计算保证效率，两者结合实现了最佳的整体可扩展性。

参数敏感性分析

M3Dphormer 的关键超参数是簇的数量 \(P\)，它影响簇掩码 \(M_{c4}\) 的质量,图展示了不同 \(P\) 值下的模型性能。总体而言，M3Dphormer 在大多数数据集上对 \(P\) 的选择表现出很强的鲁棒性。Chameleon 是一个例外，该图仅有 890 个节点，\(P\) 的变化对划分质量影响较大，导致性能波动更显著。这一结果说明簇数量的选择在小规模图上需要更谨慎的调优。

优点和讨论

个人认为，本文有如下一些优点和创新点可供参考学习：

1. 提出的统一层次化掩码框架将不同架构的图 Transformer 统一到注意力掩码设计的框架下，揭示了模型架构与掩码构建的等价性，为理解和设计图 Transformer 提供了清晰的理论基础。
2. 定理 3.1 给出的正确分类概率上下界分析为掩码设计提供了指导原则。感受野大小和标签一致性两个维度的拆解，以及对不同层次掩码在典型场景下适用性的分析有理论价值义。
3. 双层专家路由机制的设计符合层次化交互的内在结构，先确定局部专家权重再分配簇和全局专家的两级路由，相比单层门控更符合交互信息的层次关系。
4. 双注意力计算方案从理论分析出发，给出了稀疏与稠密模式切换的精确阈值条件，而非经验性的启发式规则。基于区域划分的动态模式选择充分利用了图掩码的局部稀疏性差异，在保证效率的同时实现了最优的可扩展性。

这篇论文的主要贡献在于，将图 Transformer 中反映不同颗粒度的信息的节点与掩码形式整合到一个统一的框架中，并设计合适的混合专家模型机制来对这些信息进行利用。同时这篇论文在写作上是通过理论推导来提出假设并解决问题，能体现解决问题的必要性和工作的说服力。总体来看，这项工作值得参考。