简介Graphormer 是论文《Do Transformers Really Perform Bad for Graph Representation?》的官方实现面向图神经网络与图 Transformer 方向的研究者、算法工程师及需要复现图表示学习实验的学生。它通过多种结构编码将图的拓扑信息有效注入标准 Transformer 架构在 PCQM4M-LSC、MolPCBA、MolHIV 和 ZINC 等基准上取得领先性能可用于分子性质预测等图级任务的研究与二次开发。资源包共 6 个文件包含 4 个 md 说明文档、1 个 model.py 核心模型脚本和 1 个 license 许可文件压缩后约 6KB体量轻便便于快速阅读源码与理解模型结构。目前已有 1258 人学习下载适合希望深入理解图 Transformer 编码机制、复现官方实验或在此基础上开展图表示学习研究的读者参考。1. Graphormer图神经网络里塞进 Transformer为什么反而更准了图神经网络这几年有个挺反直觉的现象很多团队把图数据硬塞进标准 Transformer效果居然比精心设计的 GNN 还好。按常理说图有拓扑结构、有边特征、有节点度数差异Transformer 那套全连接注意力机制应该水土不服才对。但 Graphormer 这篇工作恰恰证明了只要把图的结构信息用正确的方式编码进注意力Transformer 不但不差反而能在图表示学习上刷出当时的最优结果。这个结论值得搞懂因为它直接改变了我们设计图模型时的思路——不是去改注意力机制本身而是去改喂给注意力的东西。Graphormer 的官方实现就是这套思路的完整落地用 Python 写就核心代码量不大但每一处设计都值得反复琢磨。2. 结构信息怎么进注意力三个编码机制拆开看2.1 中心性编码让每个节点知道自己有多重要标准 Transformer 的输入是 token 序列每个 token 只有内容特征没有“位置”之外的任何结构信息。图数据不一样一个节点的度数可能差出几个数量级高度数节点和叶子节点在信息传播中的角色完全不同。Graphormer 的做法是给每个节点算一个中心性分数然后把这个分数映射成一个可学习的嵌入向量直接加到节点特征上。常见的实现思路是用度数作为中心性度量但度数是个无界整数直接当索引查嵌入表会爆炸。Graphormer 的处理很巧妙对度数做对数分桶把连续的度数值离散化到有限个桶里每个桶对应一个可学习的嵌入向量。这样既保留了“度数越大越重要”的单调性又控制了参数量。# 中心性编码的核心逻辑示意 import torch import torch.nn as nn class CentralityEncoding(nn.Module): def __init__(self, max_degree, hidden_dim): super().__init__() # 每个度数桶对应一个可学习向量 self.degree_emb nn.Embedding(max_degree 1, hidden_dim) self.max_degree max_degree def forward(self, node_features, in_degrees): # 度数截断超过上限的统一归到 max_degree clipped torch.clamp(in_degrees, maxself.max_degree) # 把中心性嵌入直接加到节点特征上 return node_features self.degree_emb(clipped)这段代码里max_degree是个超参通常取数据集里度数的某个分位数比如 95% 或 99%避免个别超级节点把嵌入表撑得太大。in_degrees是入度对有向图来说入度和出度可以分别编码但 Graphormer 原文只用了入度因为信息聚合主要沿着入边方向。加法的位置在进入 Transformer 层之前意味着中心性信息会参与每一层的注意力计算而不是只在第一层生效。2.2 空间编码把最短路径距离变成注意力偏置中心性编码只告诉模型“这个节点重不重要”但没告诉它“两个节点之间有多远”。图上的距离和序列上的距离完全是两码事两个节点可能在序列里挨着但在图上隔了十万八千里。Graphormer 用最短路径距离SPD来解决这个问题。具体做法是对每一对节点 (i, j)算它们之间的最短路径长度然后把这个长度映射成一个可学习的标量偏置加到注意力分数上。注意力分数原本是 query 和 key 的点积现在变成点积加上一个与 SPD 相关的偏置项。这样模型在计算注意力权重时会天然地偏好距离近的节点对。class SpatialEncoding(nn.Module): def __init__(self, max_dist, num_heads): super().__init__() # 每个注意力头有自己的一套距离偏置 self.dist_bias nn.Embedding(max_dist 1, num_heads) self.max_dist max_dist def forward(self, attn_scores, spd_matrix): # spd_matrix: [batch, num_nodes, num_nodes] clipped torch.clamp(spd_matrix, maxself.max_dist) # 取出对应的偏置形状调整为 [batch, num_heads, num_nodes, num_nodes] bias self.dist_bias(clipped).permute(0, 3, 1, 2) return attn_scores bias这里有个细节值得注意dist_bias的维度是[max_dist1, num_heads]意味着每个注意力头独立学习对距离的偏好。有的头可能关注近邻有的头可能关注远距离依赖这种多头差异化是 Transformer 表达能力强的重要原因。max_dist通常取 20 左右超过这个距离的节点对统一视为“足够远”偏置共享。SPD 矩阵的计算复杂度是 O(N^3) 的 Floyd-Warshall 或者对每个节点做 BFS对于大规模图这是个瓶颈实际实现里会预计算并缓存。2.3 边编码让注意力感知边的存在和特征前两个编码处理的是节点层面的信息但图还有边。标准注意力机制里任意两个 token 之间都有注意力权重但在图里两个节点之间可能根本没有边。Graphormer 的做法是在注意力分数上再加一个边编码偏置如果节点 i 和 j 之间有边就把这条边的特征映射成一个偏置加到注意力上如果没有边就用一个可学习的“无边”嵌入。class EdgeEncoding(nn.Module): def __init__(self, num_heads, edge_dim): super().__init__() # 把边特征投影到每个头的偏置 self.edge_proj nn.Linear(edge_dim, num_heads) # 无边情况的偏置 self.no_edge_bias nn.Parameter(torch.zeros(num_heads)) def forward(self, attn_scores, edge_features, edge_mask): # edge_features: [batch, num_nodes, num_nodes, edge_dim] # edge_mask: [batch, num_nodes, num_nodes]有边为1无边为0 edge_bias self.edge_proj(edge_features) # [batch, N, N, num_heads] edge_bias edge_bias.permute(0, 3, 1, 2) # [batch, num_heads, N, N] # 无边位置用 no_edge_bias 填充 mask edge_mask.unsqueeze(1) # [batch, 1, N, N] bias edge_bias * mask self.no_edge_bias.view(1, -1, 1, 1) * (1 - mask) return attn_scores bias边编码和空间编码是叠加的最终注意力分数 点积 空间偏置 边偏置。这种设计让模型同时感知“距离多远”和“边是什么类型”。对于分子图这种边有化学键类型的场景边编码能捕捉到单键、双键、芳香键的区别对于社交网络边编码可以编码互动频率或关系类型。注意三个编码的加法顺序不影响结果因为都是加性偏置。但边编码的no_edge_bias初始化很关键如果初始化太大模型会过度关注无边节点对实际实现里通常初始化为零或很小的值。3. 从编码到前向传播Graphormer 层的完整数据流3.1 节点特征、中心性、空间、边编码的汇合点一个 Graphormer 层的输入有四个部分节点特征矩阵、中心性编码、空间编码矩阵、边编码张量。数据流是这样的节点特征先加上中心性编码然后进入多头注意力。注意力计算时query 和 key 的点积结果上叠加空间偏置和边偏置再经过 softmax 得到注意力权重。权重作用于 value 后经过残差连接和层归一化再进前馈网络。class GraphormerLayer(nn.Module): def __init__(self, hidden_dim, num_heads, ffn_dim, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(hidden_dim, num_heads, dropoutdropout) self.norm1 nn.LayerNorm(hidden_dim) self.norm2 nn.LayerNorm(hidden_dim) self.ffn nn.Sequential( nn.Linear(hidden_dim, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, hidden_dim), nn.Dropout(dropout) ) def forward(self, x, attn_bias): # x: [N, batch, hidden_dim] # attn_bias: [batch*num_heads, N, N] residual x x self.norm1(x) attn_out, _ self.attn(x, x, x, attn_maskattn_bias) x residual attn_out residual x x self.norm2(x) x residual self.ffn(x) return x这里用的是 Pre-LN 结构即层归一化在注意力之前。Pre-LN 在深层 Transformer 里训练更稳定梯度不会因为层数增加而爆炸。attn_bias的形状是[batch*num_heads, N, N]因为 PyTorch 的MultiheadAttention要求 mask 形状和 batch 维度对齐。实际实现里会把空间偏置和边偏置合并成一个 bias 张量传进去。3.2 为什么用 Pre-LN 而不是 Post-LNPost-LN 是原始 Transformer 的做法先做注意力再加残差最后层归一化。这种结构在浅层网络里没问题但 Graphormer 通常要堆 12 层甚至更多Post-LN 会导致梯度在反向传播时经过每一层的 LayerNorm 被反复缩放深层梯度变得极小训练不动。Pre-LN 把归一化提前残差路径上的梯度可以无阻碍地回传这是深层 Transformer 能训起来的关键。另一个细节是 FFN 的维度通常是 hidden_dim 的 4 倍这是 Transformer 的标配。GELU 激活函数比 ReLU 更平滑在零点附近有非零梯度对图数据里常见的稀疏特征更友好。3.3 输出层节点级和图级任务的分叉Graphormer 的输出可以用于两类任务节点级任务如节点分类和图级任务如图分类。节点级任务直接取最后一层的节点表示接一个线性分类头。图级任务需要把整个图的节点表示聚合成一个图表示Graphormer 用一个特殊的[CLS]节点这个节点和所有其他节点都有边相连它的最终表示就是图表示。class GraphormerForGraphClassification(nn.Module): def __init__(self, num_layers, hidden_dim, num_heads, ffn_dim, num_classes): super().__init__() self.layers nn.ModuleList([ GraphormerLayer(hidden_dim, num_heads, ffn_dim) for _ in range(num_layers) ]) self.classifier nn.Linear(hidden_dim, num_classes) def forward(self, x, attn_bias, cls_mask): # x 的第一行是 [CLS] 节点 for layer in self.layers: x layer(x, attn_bias) # 取 [CLS] 节点的表示 graph_repr x[0] return self.classifier(graph_repr)[CLS]节点的引入借鉴了 BERT 的做法好处是不需要设计复杂的池化函数让模型自己学习如何聚合全局信息。cls_mask用来控制[CLS]节点和普通节点的注意力通常[CLS]可以和所有节点互相注意但普通节点之间保持原有的图结构。4. 避坑与排查Graphormer 实现里最容易翻车的几个点4.1 最短路径距离矩阵算错导致注意力偏置全乱现象模型训练 loss 不下降或者下降后验证集效果极差注意力权重可视化后发现模型在关注完全无关的节点对。原因SPD 矩阵计算时用了错误的图表示比如把有向图当无向图算或者节点索引和邻接矩阵索引对不上。更隐蔽的是对不连通图SPD 里会出现无穷大直接截断到max_dist会让所有不连通节点对的距离都一样模型无法区分“距离 20”和“根本不连通”。解决算 SPD 之前先检查图的连通性对不连通分量分别算 SPD不连通节点对用一个特殊的标记值比如max_dist1而不是直接截断。有向图必须用入边方向算 SPD因为信息沿着入边聚合。4.2 中心性编码的嵌入表太大导致显存爆炸现象训练时显存占用远超预期或者模型参数量比理论值大很多。原因max_degree设得太大比如直接用了数据集里的最大度数而某些图里存在度数上万的超级节点。嵌入表大小是(max_degree1) * hidden_dimhidden_dim 通常是 512 或 768度数上万时嵌入表就有几百万参数。解决对度数做对数分桶或者用分位数截断。常见的做法是取度数的 95% 分位数作为max_degree超过的节点统一归到最后一个桶。这样既保留了大部分节点的区分度又控制了参数量。4.3 边编码的 mask 方向搞反现象模型在无边节点对上注意力异常高有边节点对反而被忽略。原因edge_mask的定义和注意力 mask 的语义相反。注意力 mask 里 1 表示“允许注意”0 表示“屏蔽”但边编码里edge_mask是 1 表示“有边”0 表示“无边”。如果直接把edge_mask当注意力 mask 用就会把有边的地方屏蔽掉。解决明确区分两个 mask 的语义。边编码的 mask 只用来选择用边特征还是无边嵌入不参与注意力的屏蔽。注意力屏蔽用单独的attn_mask通常只在 padding 位置屏蔽。4.4 预计算的 SPD 矩阵没有缓存导致训练极慢现象每个 epoch 开始前有很长的停顿GPU 利用率在数据加载阶段几乎为零。原因SPD 矩阵在__getitem__里现算每个 batch 都要跑一遍 BFS 或 Floyd-Warshall。对于节点数几百的图单次计算可能只要几毫秒但乘以几万个图再乘以 epoch 数累积起来就是几小时。解决在数据集初始化时一次性算好所有图的 SPD 矩阵存成内存映射文件或直接放内存。如果内存不够用 LRU 缓存最近用过的图。SPD 矩阵只和图的拓扑有关和节点特征无关所以可以跨 epoch 复用。4.5 多头注意力的 bias 维度对不上现象运行时报维度错误或者模型能跑但效果和单头一样。原因attn_bias的形状要求是[batch*num_heads, N, N]但空间编码和边编码输出的形状是[batch, num_heads, N, N]。如果直接相加后传给MultiheadAttentionPyTorch 会按广播规则处理结果可能不是预期的。解决在传给注意力之前把 bias 的形状从[batch, num_heads, N, N]reshape 成[batch*num_heads, N, N]。用permute和reshape组合注意permute后张量不连续reshape会报错需要先.contiguous()。5. 设计取舍为什么 Graphormer 没有走稀疏注意力的路Graphormer 最让人纠结的地方是它的注意力复杂度O(N^2) 的注意力矩阵加上 O(N^2) 的空间偏置和边偏置对于节点数上千的图显存和计算量都吃不消。很多图 Transformer 的后续工作都在做稀疏注意力比如只关注 k 跳邻居或者用线性注意力近似。但 Graphormer 原文坚持了全连接注意力只在编码层面做文章。这个取舍背后的逻辑是图的结构信息已经通过空间编码和边编码注入到了注意力偏置里全连接注意力让模型自己决定哪些远距离节点对值得关注。如果强行稀疏化可能会切断一些重要的长程依赖。比如在分子图里两个官能团可能隔了很多化学键但它们之间的相互作用决定了分子的性质稀疏注意力很容易漏掉这种关系。从实现角度看Graphormer 的官方代码在效率上做了不少优化。SPD 矩阵预计算并缓存边编码用矩阵乘法批量处理中心性编码的嵌入表用分桶控制大小。这些优化让它在中等规模图上节点数几百到一千能跑得动但再大就得上采样或分块了。我一般会这么理解 Graphormer 的定位它不是要替代所有 GNN而是在图结构信息丰富、节点数不太大、需要捕捉长程依赖的场景下提供了一个比传统 GNN 更强的选择。如果你的图有几十万节点Graphormer 的全连接注意力确实不现实但它的编码思路——中心性、空间、边——可以移植到任何图 Transformer 里哪怕你用的是稀疏注意力。从那以后我每次设计图模型都会先问自己三个问题节点的重要性差异大不大节点间的距离信息重不重要边有没有特征如果三个答案都是“是”那 Graphormer 的编码方案就值得认真考虑。这套官方实现把三个编码机制拆得很清楚代码结构也干净适合拿来当图 Transformer 的参考实现改起来比从头写省事得多。希望帮到你。本文还有配套的精品资源点击获取