Graphormer结构编码：图Transformer表达力与工程避坑-育师

简介：Graphormer 官方实现代码包面向图表示学习、图神经网络研究者，以及希望复现 Transformer 在图任务上表现的中高级开发者，核心解决如何把节点连接与图结构信息有效编码进标准 Transformer 的问题。包内共 6 个文件，以 4 个 Markdown 文档、1 个 Python 模型脚本和 1 个许可证文件为主，压缩包约 6KB，体量轻便，适合快速阅读模型定义、了解项目说明与使用约束。内容围绕 Graphormer 的结构编码思路展开，作者在 PCQM4M-LSC 上报告 0.1234 MAE、MolPCBA 31.39 AP、MolHIV 80.51 AUC、ZINC 0.122 MAE 等结果，可为图性质预测、分子建模与基准复现提供参考。已有 1258 人学习，适合作为图 Transformer 入门研读与二次开发的起点。

1. 结论先摆出来：图上的 Transformer 不是“不行”，是没把结构喂进去

在 Python 圈子里讨论 graph transformer 时，最常被问的反问是：把自注意力原样套到图上，为什么经常打不过一个几百行的消息传递网络？问题不在于 Transformer 本身弱，而在于图没有天然的序列顺序：节点编号一换，注意力看到的东西就全变了。Graphormer 的官方实现给出的解法很朴素：不改注意力，改注意力的输入和偏置，用三种结构编码把“谁离谁多远、谁连了谁、谁的度有多大”直接写进注意力分数里。搞懂这三处注入点，基本就搞懂了整类图 Transformer 的骨架。

2. 自注意力对结构“失明”：置换不变性到底丢掉了什么

把一张图输入标准 Transformer，做法通常是把节点特征排成序列、加位置编码，然后跑多头注意力。问题在于序列有先后，图没有：任何节点编号的置换都不该改变输出，这叫置换不变（图级任务）或置换等变（节点级任务）。自注意力天然满足这一点，因为 QK 内积只看两个向量，不看它们的下标。但“对称”在这里是双刃剑：它同时对“这两个节点之间到底有没有边”也一视同仁。对一张 10 个节点的图，注意力看到的只是一个 10×10 的全连接打分矩阵，边是否存在这层最关键的信息，被彻底抹平了。

很多实现的第一反应是补一个节点位置编码，比如按度排序发一个 index，或者跑一遍随机游走拿序列位置。这条路能缓解，但补进去的是“节点在某个虚构序列里的位置”，不是“节点和节点之间的拓扑关系”。前者是每个节点一个向量，后者本质上是每对节点一个量，维度都不一样，硬塞进节点特征里必然丢信息。Graphormer 的出发点就在这里：结构信息不是节点属性，是节点对属性，那它就该待在对属性该待的位置——注意力分数上。

2.1 消息传递网络的天花板：1-WL 测试这条线

要理解为什么非要动注意力，得先看清原有方案的边界。主流图神经网络走的是消息传递：每一层把邻居特征聚合过来更新自己，堆 L 层就能看到 L 跳以内的邻居。这条路的表达能力有个著名的上界，和 1-WL 图同构测试等价——凡是一层邻居聚合区分不开的两张图，堆再多层也区分不开。

最经典的反例是六元环和两个三元环拼起来的图。两张图节点数相同、每个节点的度都是 2、连边数相同，1-WL 迭代多少轮给出的颜色分布都一模一样，于是消息传递网络给出的图表示也几乎相同。但这两张图显然不同构：一个连通，一个不连通；一个任意两点最远距离是 3，另一个跨团两点的距离分布完全不同。

提示：这不是实现 bug，是机制本身的信息瓶颈。聚合算的是“邻居特征的和或均值”，它能数局部度数，数不出全局的距离分布，也算不出环的长度。

2.2 全连接注意力反而是一种“超纲”的表达力

Graphormer 的选择很干脆：既然一跳一跳聚合看不见全局，那就干脆让每个节点直接看到所有其他节点，把图当成完全图输入。注意力本身在完全图上跑，但额外给每对节点 (i, j) 一个可学习的偏置项，这个偏置由它们之间的结构关系决定。上面那个六元环的反例，在“任意两点最短路径距离”这个视角下立刻可分：六元环的距离分布里没有距离 1 之外的跨团边，两团图里存在大量距离 2 的跨团点对，偏置表会把这两种分布映射到不同的分数上。

这里有个非常关键的设计判断，也是最容易做错的一步：把结构做成 mask 还是做成 bias。

# 方式 A：邻接矩阵当 mask，只允许邻居互相看见 # 等价于注意力版的 GAT，感受野仍然是一跳，表达力回到 1-WL 以内 attn = attn.masked_fill(adj == 0, -1e9) # 方式 B：所有节点对都可见，把结构做成加性偏置 # 感受野是全局，结构作为先验参与打分，才可能突破 1-WL attn = attn + spatial_bias + edge_bias

方式 A 写起来最自然，也最符合“图就是稀疏的”这个直觉，但它把 Transformer 退化成了消息传递，白搭了全连接注意力的表达力。方式 B 才是 Graphormer 的路线：mask 只用来处理 padding，结构关系一律走加性偏置。这个差别在代码上就是两个字符，在效果上是一整条表达力上界的分界线。代价也很直接——注意力矩阵是 N×N 的稠密张量，图一大显存就撑不住，这笔账后面单独算。

3. 三种结构编码：中心性、空间、边，分别补哪一块信息

结构信息至少有三个层次：节点自己在局部有多重要（度）、任意两点在拓扑上隔多远（距离）、以及这条路径上具体的边长什么样（边特征）。三者维度不同、注入位置也不同，Graphormer 用三套编码分别对应，最终合成一个加在 softmax 之前的偏置张量。完整打分公式大致长这样：

A(i,j) = (h_i W_Q)(h_j W_K)^T / sqrt(d) + b(φ(i,j)) + c(i,j)

前一项是标准注意力，b 是按最短路径距离取的空间偏置，c 是沿路径累积的边编码。节点特征侧还有一份中心性编码，加在输入上。四项各管一段，缺一项消融就会掉点。

3.1 中心性编码：把度数加进节点特征

度是节点最廉价也最有信息量的属性。分子图里度过大往往意味着这是个连接枢纽原子，社交图里高度节点是枢纽用户。标准注意力对这件事一无所知，因为它只看特征内积。中心性编码的做法极其简单：维护两张 embedding 表，按入度和出度查表，把查到的向量直接加到节点特征上。

import torch import torch.nn as nn class CentralityEncoding(nn.Module): def __init__(self, hidden_dim, max_degree=512): super().__init__() # 入度 / 出度各一张表；有向图上两者语义不同，必须分开 self.deg_in = nn.Embedding(max_degree + 1, hidden_dim) self.deg_out = nn.Embedding(max_degree + 1, hidden_dim) def forward(self, x, in_degree, out_degree): # x: [B, N, D] 节点特征（如原子类型 embedding） # in_degree / out_degree: [B, N] 的 long 张量 cap_in = self.deg_in.num_embeddings - 1 cap_out = self.deg_out.num_embeddings - 1 return (x + self.deg_in(in_degree.clamp(max=cap_in)) + self.deg_out(out_degree.clamp(max=cap_out)))

参数上只有两个关键项：max_degree决定桶的数量，hidden_dim必须和节点特征维度一致，因为这里是逐元素相加而不是拼接。相加的好处是不增加序列长度、不改变后续所有张量的形状；拼接会让隐藏维度翻倍，后面每一层都要跟着改。为什么用 embedding 而不是 one-hot 过一层线性层？数学上两者等价，但 embedding 写法在超大度数表上更省显存，梯度也只在被命中的那几行上更新。

clamp那一行是整个模块唯一的坑源。真实图的度分布是长尾的，少数节点度能到几千，为它们单独开桶既浪费又过拟合，所以普遍做法是截断到一个上限，超出的全部共用最后一个桶。这是一种有意的信息压缩：度 800 和度 5000 的节点拿到同一个向量。截断阈值该按数据集的度分位数来定，拍脑袋设 512 在大多数分子图上够用，在社交图上大概率不够。

3.2 空间编码：把最短路径距离变成每头独立的偏置标量

空间编码管的是“两点隔多远”。预处理阶段先算全图的最短路径距离矩阵 φ，不可达的点对单独标记。然后为每一种距离准备一个可学习标量，加在对应的注意力 logits 上。

class SpatialEncoding(nn.Module): def __init__(self, num_heads, max_dist=32): super().__init__() # 每个 head 对每种距离有独立标量；max_dist 之外合并成一个桶， # 不可达再单独占一个桶，所以是 max_dist + 2 self.bias = nn.Parameter(torch.zeros(num_heads, max_dist + 2)) self.max_dist = max_dist def forward(self, spd): # spd: [B, N, N] 最短路径长度，-1 表示不可达 idx = spd.clone() idx[idx < 0] = self.max_dist + 1 # 不可达 -> 专门的桶 idx = idx.clamp(max=self.max_dist) # 超过上限的距离合并 # self.bias[:, idx] 得到 [H, B, N, N]，整理成 [B*H, N, N] b = self.bias[:, idx] b = b.permute(1, 0, 2, 3).reshape(-1, idx.size(1), idx.size(2)) return b

参数量是num_heads × (max_dist + 2)，小到可以忽略。这个“每头一个标量”的选择值得多说一句：如果改成每头一个向量再和某个投影做内积，参数量会涨到heads × D²，收益却很小，因为距离本身只是个有序的离散量，一个标量足以表达“近的加分、远的减分”。更重要的是，加在 logits 上等价于给注意力权重乘一个 exp(b) 的因子，是一种乘性调节；如果把距离拼进节点特征再走 QK 内积，效果会被特征内积本身冲淡，等于绕了一圈才回到原点。

多头在这里的作用也不是摆设。不同的头可以学出完全不同的距离偏好：有的头只关注一跳邻居（距离 1 的偏置远高于其他），有的头做长程聚合（距离大的偏置反而更高），有的头对不可达点对给极低分。这种分工在训练后可以直接把bias打出来看，是排查模型到底学没学到结构的一个很实用的手段。

3.3 边编码：把路径上的边特征压成一个标量

空间编码只知道“隔了几跳”，不知道“这几跳上的边是什么样”。分子图里键的类型（单键、双键、芳香键）对性质影响极大，两个原子间距离都是 2，中间经过的是单键还是芳香键，化学意义完全不同。边编码补的就是这一块。

它的做法是：对每对节点 (i, j)，取出最短路径上经过的所有边特征，每一跳配一个可学习权重向量，做加权内积后对跳数取平均，得到一个标量加进注意力分数。

class EdgeEncoding(nn.Module): def __init__(self, num_heads, edge_dim, max_hop=8): super().__init__() # 每一跳一个权重向量，作用在该跳的边特征上 self.hop_weight = nn.Parameter(torch.zeros(num_heads, max_hop, edge_dim)) def forward(self, edge_feat): # edge_feat: [B, N, N, max_hop, E] # 预处理阶段已按最短路把每对节点的第 k 跳边特征排好，无该跳则填 0 w = self.hop_weight.mean(dim=0) # [max_hop, E] score = (edge_feat * w).sum(-1).mean(-1) # 先对 E 加权求和，再对跳数平均 return score # [B, N, N]

参数说明：edge_dim是原始边特征维度，max_hop是路径截断长度，num_heads让每个头有自己的一套跳权重。这里对跳数取平均而不是求和，是为了让不同长度路径的编码落在同一量级上，否则长路径天然占优，模型会偏向于关注远距离点对。无该跳的位置填 0 再平均，等于把短路径的贡献按实际长度归一化。

代价在预处理：要为所有点对存下max_hop跳的边特征，张量形状是 N×N×max_hop×E。这个张量通常是整条流水线里最占内存的东西，也是后面要讲的显存坑的主要来源。三种编码最终在 batch 维度上拼成统一的attn_bias，一次加进注意力 logits。

4. 一次前向到底发生了什么：从节点特征到 [CLS] 读出

三种编码各自就位后，还得看它们怎么被装配进完整的网络。这一层最容易出问题的地方不在数学，而在工程顺序：谁是 pre-LN、谁是 post-LN、mask 加在哪一步、读出取哪个位置的向量，每一处写错都能让训练悄悄变差而不报错。

4.1 输入侧：节点特征、虚拟节点与 mask

节点特征通常是查表得到的 embedding（原子类型、节点类型），加上中心性编码后进入主干。除此之外还要额外插入一个特殊节点，习惯上叫它 [CLS] 或虚拟节点：它和所有真实节点之间都有连接，本身的特征是可学习的。这个节点的作用等价于 GNN 里的全局节点，但比全局池化更灵活——池化是把所有节点向量做平均或求和，信息在池化那一刻被压扁；虚拟节点则是通过注意力逐层从所有节点收集信息，每一层都能选择性地取，最后直接拿它的输出做图级表示。

mask 在这一步就要准备好。一个 batch 里的图节点数不同，短图要 pad 到最长的那张，pad 出来的位置在注意力和损失里都必须屏蔽。注意力和损失用的是两套 mask：注意力用 key padding mask 屏蔽列，损失用节点级 mask 只对真实节点算。两处漏掉任何一处，梯度都会被 padding 位置的无效计算污染。

4.2 注意力层的拼装顺序与 pre-LN 的稳定性

主干的每个 block 由两层子层构成，顺序是“归一化 → 子层 → 残差”，也就是 pre-LN。这个选择不是风格问题：post-LN 在深层 Transformer 上训练初期梯度波动很大，需要 warmup 才能稳住；pre-LN 把 LayerNorm 挪到子层入口，残差路径上是一条干净的恒等映射，梯度可以无衰减地回传，深层堆叠时明显更稳。

import math import torch import torch.nn as nn class GraphormerAttention(nn.Module): def __init__(self, hidden_dim, num_heads): super().__init__() self.num_heads = num_heads self.head_dim = hidden_dim // num_heads self.qkv = nn.Linear(hidden_dim, hidden_dim * 3) self.out = nn.Linear(hidden_dim, hidden_dim) def forward(self, x, attn_bias, key_padding_mask=None): B, N, D = x.shape q, k, v = self.qkv(x).chunk(3, dim=-1) # 拆头：[B, N, H, d] -> [B, H, N, d] q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(B, N, self.num_heads, self.head_dim).transpose(1, 2) logits = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) # attn_bias: [B*H, N, N]，空间偏置与边偏置之和，加在 softmax 之前 logits = logits + attn_bias.view(B, self.num_heads, N, N) if key_padding_mask is not None: # [B, N] -> [B, 1, 1, N]，只屏蔽被 pad 的列 logits = logits.masked_fill(key_padding_mask[:, None, None, :], -1e9) attn = logits.softmax(dim=-1) out = torch.matmul(attn, v).transpose(1, 2).reshape(B, N, D) return self.out(out) class GraphormerBlock(nn.Module): def __init__(self, hidden_dim, num_heads, ffn_dim, dropout=0.1): super().__init__() self.ln1 = nn.LayerNorm(hidden_dim) # pre-LN：先归一化再进子层 self.attn = GraphormerAttention(hidden_dim, num_heads) self.ln2 = 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), ) self.dropout = nn.Dropout(dropout) def forward(self, x, attn_bias, key_padding_mask=None): x = x + self.dropout(self.attn(self.ln1(x), attn_bias, key_padding_mask)) x = x + self.dropout(self.ffn(self.ln2(x))) return x

三个细节值得单独拎出来。第一，attn_bias是在 softmax 之前加的，这是它能起作用的唯一位置，加到 softmax 之后等于对权重做线性平移，信息全丢。第二，mask 用的是 -1e9 而不是-inf，因为整行全被 mask 时-inf的 softmax 会算出 0/0 得到 NaN，一个大负数在数值上等价于屏蔽又不会炸。第三，堆完所有 block 之后还要在最外层再套一个 LayerNorm，这是 pre-LN 结构里经常被漏掉的一步——残差路径上从未被归一化过，输出尺度会随层数漂移，补上这一层之后接分类头才稳。

4.3 读出与任务头

读出方式取决于任务粒度，这一点在实现里必须分清楚。图级任务（图分类、图回归、分子性质预测）直接取 [CLS] 节点在最后一层的向量，接一个两层 MLP 输出；节点级任务（节点分类）取真实节点的向量，同样接 MLP，只是每个节点独立出结果；边级任务则在节点表示的基础上再做一次点对组合。三种读出头共享主干，只是取的位置不同。

值得一提的是，有些实现会额外加一份绝对位置编码或 3D 坐标编码，这属于可选增强而不是核心机制。真正让 Graphormer 和“Transformer 加个位置编码”区分开的，是那三种基于图结构的编码，它们在预训练权重里也是被重点训练的部分。

5. 避坑与排查：这几种翻车几乎每个实现者都会遇到

结构编码的逻辑只有几十行，但把它跑对、跑稳、跑到和报告里的数字一致，中间隔着不少血泪经验。下面五条是重写实现时最常撞上的，按“现象 → 原因 → 解决”排一遍。

5.1 结构编码里的三个隐形错误

度数截断把枢纽节点全压成一个向量。现象是整体 loss 曲线正常，但在少数高度节点的样本上预测出现系统性偏差，验证集某些子图上的误差明显高于平均。原因是中心性编码的clamp把超过阈值的度全部映射到同一个桶，模型完全无法区分度 600 和度 3000 的节点。解决办法是先统计训练集的度分布，按 95 或 99 分位定阈值，或者改用对数分桶——把度取 log 后再离散化，长尾部分依然能保留区分度。

“不可达”和“距离 0”混进了同一个桶。现象是在分子图这种多连通分量的小图上精度莫名下降，消融实验里去掉空间编码反而没掉多少，说明这份编码没学到东西。原因是预处理时用 0 表示不可达，而 0 同时也是对角线的距离值，编码表里这两类点对被迫共享一个标量，模型于是学会忽略它。解决办法是给不可达单独分配一个索引（就像前面代码里的max_dist + 1），并在生成 mask 时保持一致。

边特征张量的跳数对齐错位。现象是模型能训练但收敛极慢，或者去掉边编码后效果几乎不变。原因是预处理阶段回溯最短路径时，第 k 跳的边特征被写进了错误的槽位，导致权重向量学到的是一堆错位的统计量。这类 bug 不会报错，只会让特征变成噪声。排查手段是构造一个只有三四个节点的小图，手工算出每对节点的路径，逐跳打印出来比对。

5.2 张量与数值：显存和 NaN 的两个坑

padding 位置参与 softmax 导致 loss 变 NaN。现象是训练跑了几十步之后 loss 突然变成 nan，之后再也回不来。原因通常有两个叠加：mask 用了-inf，而某张图的某一整行全部被屏蔽（比如 [CLS] 之外的孤立节点），softmax 分母为 0；或者损失函数没有按节点 mask 加权，padding 位置的梯度污染了参数。解决办法是把 mask 值换成 -1e9 这类大负数，并在算 loss 时用loss * node_mask再归一化。

N×N 的偏置张量直接把显存吃满。现象是节点数超过几百就 OOM，而模型参数量明明很小。原因是空间偏置、边编码、注意力 logits 三份张量都是 O(N²) 级别，边编码还多乘一个max_hop × edge_dim。这块没有魔法，只有取舍：把 batch 内按节点数排序做长度分桶，减少 padding 浪费；边编码的max_hop从 8 降到 4；超大图改成子图采样训练。如果一张图动辄上万节点，那这套全连接注意力本身就不合适，该换稀疏方案，而不是硬扛。

6. 表达力与成本的那笔账：为什么它敢用全连接的 O(N²) 注意力

三种编码加上全连接注意力，本质上是用计算复杂度换表达力。消息传递网络的时间复杂度大致和边数线性相关，稀疏图上极其便宜，但表达力被 1-WL 卡死；全连接图 Transformer 把每对节点都算一遍，复杂度顶到 O(N²·D)，换来的是任意两点之间的结构关系都能直接参与打分，不再受跳数限制。中间那条路——稀疏注意力图 Transformer——只让部分点对可见，省了算力，但也把一部分长程关系挡在门外，表达力介于两者之间。

方案	结构信息进入方式	感受野	表达力上界	复杂度
消息传递网络	邻居聚合	一跳/层	受 1-WL 限制	约 O(N·E·D)
稀疏注意力图 Transformer	邻接 mask + 局部偏置	受稀疏模式限制	介于两者之间	与可见边数相关
全连接 + 结构偏置	加性偏置注入 logits	全局	可超越 1-WL	O(N²·D)

这张表也解释了它的适用边界：中小规模图（分子、小分子晶体、几百节点以内的图级任务）是这个设计的甜区，O(N²) 完全承受得起，表达力收益却是实打实的；动辄十万节点的社交图或引文图，这套东西根本跑不动，得回到采样或稀疏路线。它不是“更快的 GNN”，而是“在算力允许的范围内把表达力顶到最高”的一种取舍。

验证这套机制有没有真的生效，我一般会做三步。先在小图上暴力核对：手写一个 Floyd 求最短路，和预处理产出的距离矩阵逐元素比对，确保偏置索引没错位。再做单项消融：分别去掉中心性、空间、边编码各跑一遍，看验证指标掉多少，掉得不明显的那个模块大概率有 bug 而不是没用。最后把训练好的bias参数打出来看分布，如果所有 head 学到的距离偏好几乎一样，说明多头没分工，可能是attn_bias在拆头那一步广播错了。