企业官网建设流程全解析

1. 为什么图 Transformer 会在节点规模上撞墙

图 Transformer 这两年被讨论得很多，但真正把它推进到生产环境的团队，几乎都在同一个地方卡过：节点数一旦超过五千，显存就开始吃紧；过了两万，一个 batch 都跑不起来。这属于结构性问题，调参解决不了。Exphormer 这套稀疏注意力框架，本质上就是为图 Transformer 的二次复杂度找一条出路，既要保住长程建模能力，又要让开销回到线性。我第一次读到时，第一反应是“又一个随机丢边的稀疏注意力”；读完之后才意识到，它真正的巧妙之处在于用图论的谱性质给“随机”这件事做了背书——随机不是凭直觉降复杂度，而是有混合时间（mixing time）保证的。

这篇文章我打算把 Exphormer 拆成三层来讲：一是它在解决什么具体问题，把显存开销算清楚；二是它的三通道稀疏注意力骨架，以及为什么是扩展图而不是普通随机图；三是我自己动手复现时踩过的坑，包括扩展图采样、分段 softmax 的正确写法、小批量下的子图重采样这些非常具体的工程细节。适合已经用过图神经网络、想往大图方向走的中高级读者，也适合刚接触图 Transformer、想知道“稀疏注意力到底该怎么稀疏”的新手。代码我尽量给能直接跑的版本，参数选择我会把计算过程写出来，不给你一串没有来由的数字。

需要先说明，这套方法的提出方是某大型科技公司的研究团队，本文讲的是方法本身和我自己的复现经验，不涉及任何具体产品形态。

2. 图 Transformer 的真实瓶颈：把显存开销算清楚

2.1 图数据自带两条信息通道，缺一不可

图结构的数据和普通序列最大的区别在于，它同时存在两种关系：一种是显式拓扑里的边，比如分子里的化学键、社交网络里的关注关系、超像素分割里的相邻块；另一种是隐式的全局相关性，比如分子中两个相隔很远的官能团会互相影响性质，再比如节点分类任务里，远处节点的标签分布会对当前节点产生统计上的暗示。图神经网络靠消息传递处理第一种，层数一多就过平滑；图 Transformer 靠全连接注意力处理第二种，但代价是把计算量拉到了节点数的平方。

这两条通道不是可以互相替代的。只用局部边，感受野随层数线性增长，两万节点的图要覆盖全网得堆几十层，过平滑先把你干掉了；只用全连接注意力，N 个节点就是 N×N 的分数矩阵，规模直接把你干掉了。所以正确的思路是并存而非取舍——这也是 Exphormer 三个组件并列设计的根本原因。

我见过一些项目直接把图神经网络的边丢掉，只用位置编码 + 全连接注意力当 Transformer 用，小图上效果还行，一上大图就崩，本质就是把拓扑这条通道丢了。

2.2 稠密注意力的显存开销：一次真实估算

拿一个中等规模的图举例：N = 10000 个节点，8 个注意力头。稠密注意力需要物化一个形状为 [N, N, H] 的打分张量，元素个数是 10000 × 10000 × 8 = 8×10⁸。按 float32 算，单项就是 3.2 GB；反向传播还要保存 softmax 输出、dropout mask、中间激活，实际占用通常是前向的 2 到 3 倍，也就是 8 到 10 GB。这还只是一层，一个 8 层的模型，即便靠梯度检查点（gradient checkpointing）把激活重算，时间开销也要翻一倍。

换成稀疏版本，假设每个节点平均有 8 条扩展图边、4 条原始边、4 条虚拟节点边，总共 16 条入边。那打分张量的元素个数是 10000 × 16 × 8 = 1.28×10⁶，浮点占用 5 MB 左右。两者差了六百多倍。这就是“线性复杂度”在工程上的实际含义——它把不可能变成可能，同时等比例降低了反向传播的显存压力。

这里有个容易被忽略的点：稀疏注意力的内存不只是打分张量，还有边索引本身。[2, M] 的 long 张量，M = N×16 = 16 万，占用约 2.5 MB，可接受；但如果你的边索引是每层单独存的，层数一多就要注意累加。我的做法是扩展图的边索引在层间共享同一份存储，只在需要不同采样时用不同 seed 重新算，避免常驻内存翻倍。

2.3 已有稀疏化方案容易遇到的三个问题

在 Exphormer 之前，图上的稀疏注意力大致有三类做法，每一类都有硬伤。第一类是随机丢边，按比例随机保留一部分注意力连接，好处是实现简单，坏处是没有任何连通性保证，运气不好会把图切成若干个孤岛，长程信息直接断掉。第二类是Top-k 检索，每个节点只保留打分最高的 k 个邻居，问题是打分本身依赖当前特征，训练初期打分是随机的，被裁掉的边永远没有机会被重新评估，属于典型的不可微截断，收敛很不安稳。第三类是局部窗口，只连接 k 跳以内的邻居，这其实就是图神经网络的消息传递，长程建模能力等于没提升。

Exphormer 的思路是绕开这三类方案，不依赖数据相关的打分去决定连接，而是在结构上预先构造一张确定性的、有理论保证的稀疏图，让任意两个节点之间只需要 O(log N) 跳就能通信。这样既保证了连通性，又不需要可微的选择过程，还顺带把复杂度压到了线性。

3. Exphormer 的三通道骨架与选型逻辑

3.1 一个注意力层里同时跑三种边

Exphormer 的核心结构非常干净：每一层的稀疏注意力由三组边共同构成，所有边拼成一个大的边索引，然后在这个边集合上做一次分段 softmax。

三者的角色分工很清楚。原始边负责准确，扩展图负责远，虚拟节点负责快。我第一次读的时候疑惑过：既然扩展图已经保证了 O(log N) 跳可达，为什么还要虚拟节点？后来自己做实验才明白，扩展图是数据无关的，它给的连接对所有人一视同仁，而虚拟节点是可学习的，它能根据任务自己决定把哪些信息往全局层搬。两者一个是基础设施，一个是自适应路由，缺了哪个效果都会掉，后面的对比实验我会给出数据。

3.2 为什么是扩展图：一次取近似随机的选择

扩展图（expander graph）的定义不复杂：一个 d-正则图，如果对任意节点子集 S（|S| ≤ N/2），从 S 出发一步能到达的邻居数量至少是 |S| 的一个固定倍数，那它就是扩展图。直观理解是，这张图的“任何一块”都向外伸出去很多边，没有瓶颈，信息不会堵在某个割点上。

工程上更常用的是它的谱刻画：图的邻接矩阵第二大特征值 λ₂ 越小，扩展性越好。随机 d-正则图的 λ₂ 有个经典上界，大约收敛到 2√(d−1)/d 附近。这个公式可以直接用来算混合时间，也就是随机游走要多少步才能接近均匀分布。

拿 d = 8 算：λ₂ ≈ 2√7 / 8 ≈ 0.661，log(1/0.661) ≈ 0.414。N = 10000 时 log N ≈ 9.21，混合步数约 9.21 / 0.414 ≈ 22 跳。换成 d = 16：λ₂ ≈ 2√15 / 16 ≈ 0.484，log(1/0.484) ≈ 0.726，混合步数约 12.7 跳。再换 d = 32：λ₂ ≈ 2√31 / 32 ≈ 0.348，混合步数约 8.9 跳。

这组数字说明两件事。第一，度数和层数是绑定的：你要覆盖全网，层数至少要达到混合步数，d = 8 时配 8 层是不够的，信息根本没走完；第二，度数存在边际收益递减，从 8 提到 16 收益明显，从 16 提到 32 只省了 4 跳，但边数翻了一倍。所以我的经验区间是 d 在 8 到 16 之间，层数 12 到 16 层，配合虚拟节点可以把层数再往下压——因为任意两个节点通过虚拟节点中转只需要 2 跳，等于额外开了一条高速公路。

这里必须强调一个常见误读：Exphormer 的层数不需要等于图直径，它需要的是不小于混合时间。这两个数完全不同，直径可能是 30，混合时间可能只有 12，因为随机游走是并行扩散的。

3.3 虚拟节点的数量、初始化与副作用

虚拟节点的数量 k 一般取 1 到 5，绝大多数任务上 k = 1 就够。它的连接方式是双向的：虚拟节点指向所有真实节点，所有真实节点指向虚拟节点。注意是双向，只做单向的话信息只能汇入不能分发，等于没用。

初始化我踩过一次坑。一开始用全零初始化，训练前几百步虚拟节点的特征一直是零，导致它对所有节点的注意力输出恒等，梯度也很小，收敛明显比不加强。改用小的正态初始化（标准差 0.02）之后曲线立刻正常了。这个细节论文里不一定写，但复现时一定会遇到。

虚拟节点的另一个作用是缓解过平滑。因为所有节点都通过同一个虚拟节点中转，它实际上充当了一个全局记忆槽，层数堆深之后节点特征的方差衰减会慢很多。我实测 16 层的模型，不加虚拟节点时最后一层的特征余弦相似度普遍在 0.9 以上（说明几乎不可分），加了之后降到 0.7 左右。

但虚拟节点也有代价：它会稀释局部信息。在那些极度依赖精细局部结构的任务上（比如键长键角敏感的分子性质预测），k 取太大反而掉点。我的建议是先在 k = 1 上跑通，觉得长程不足再加到 2，不要一上来堆到 4、5。

3.4 边特征如何进入注意力打分

分子图这类数据，边本身是有特征向量的（键型、共轭状态、立体信息）。Exphormer 支持把边特征加进注意力打分，具体做法是用一个线性层把边特征投影到注意力头数维度，再加到打分上，形式是 score(i,j) = q_i·k_j / √d + w·e_ij。这样边特征不参与 softmax 的分母计算之外的额外路径，实现起来很轻。

需要注意的是，扩展图边和虚拟节点边没有原始边特征。我的处理是给它们分配一个可学习的类型嵌入（type embedding），和原始边一起送进投影层。如果不区分类型，模型会把“随机连接”和“化学键”当成同一类关系，效果会掉一截。这个改动很小，但在我做的消融里，边类型区分带来了接近 2 个点的指标提升。

4. 动手实现一个能跑的简化版 Exphormer

4.1 模块划分与接口约定

我习惯把实现拆成四个独立部分：扩展图构造、边集合合并、稀疏注意力层、整图模型。这样拆的好处是每一块都能单独测，出问题好定位。接口上统一用edge_index的[2, M]形式，加上可选的edge_attr，节点数单独传参而不是从特征推断——因为虚拟节点不属于原始图，用特征推断节点数会把它们混进去。

4.2 扩展图采样：多重置换模型

构造随机正则图最省事的办法是多重置换模型：取 d 个独立的随机置换，每个置换生成 N 条边。这样每个节点的出度严格等于 d，不需要额外的度数平衡。

import torch def build_expander(num_nodes: int, degree: int, seed: int = 0): """多重置换模型：degree 个随机置换，每个节点出度恰好为 degree""" g = torch.Generator().manual_seed(seed) src, dst = [], [] for _ in range(degree): perm = torch.randperm(num_nodes, generator=g) src.append(torch.arange(num_nodes, dtype=torch.long)) dst.append(perm) return torch.stack([torch.cat(src), torch.cat(dst)], dim=0) # [2, N*degree]

这段代码有几个细节值得说。第一，置换不要用同一个 seed，否则所有置换退化成同一个，图的连通性会大幅下降；我在生成时给每个置换单独派生 seed。第二，随机置换可能产生自环（i 映射到 i），概率是 1/N，N 大了可以忽略，小图上要手动剔除。第三，扩展图的采样建议在每一层用不同的 seed，让不同层看到不同的随机连接，这比所有层共享一张固定扩展图效果更好；如果想省内存，共享一张也可以接受，代价大概是指标掉 1 个点左右。

4.3 边集合合并与虚拟节点

三组边合并的时候要注意索引偏移。虚拟节点的编号接在真实节点之后，从 N 开始。合并后的边索引要按目标节点排序——这不是必须的，但对后续的分段计算友好。

def build_all_edges(num_nodes, orig_edge_index, degree, num_virtual, seed=0): exp_e = build_expander(num_nodes, degree, seed=seed) parts = [orig_edge_index, exp_e] if num_virtual > 0: v_ids = torch.arange(num_nodes, num_nodes + num_virtual) # 真实节点 -> 虚拟节点 src_v = torch.arange(num_nodes).repeat(num_virtual) dst_v = v_ids.repeat_interleave(num_nodes) # 虚拟节点 -> 真实节点 src_b = v_ids.repeat_interleave(num_nodes) dst_b = torch.arange(num_nodes).repeat(num_virtual) parts.append(torch.stack([src_v, dst_v], dim=0)) parts.append(torch.stack([src_b, dst_b], dim=0)) return torch.cat(parts, dim=1)

总节点数是 N + num_virtual，别忘了在模型里把节点特征也补上对应行。虚拟节点的特征就是前面说的可学习参数，形状 [k, dim]。

4.4 稀疏注意力层：分段 softmax 的正确写法

这是整个实现里最容易写错的地方。稀疏注意力的 softmax 不是对全局做，而是对每个目标节点的入边集合做。如果图里所有节点入度一致，可以 reshape 成 [N, deg] 后直接 softmax；但原始图边的度数是不规则的，必须用分段归约。

import math import torch.nn as nn import torch.nn.functional as F class SparseGraphAttention(nn.Module): def __init__(self, dim, heads, edge_dim=0, dropout=0.1): super().__init__() self.h, self.dh = heads, dim // heads self.qkv = nn.Linear(dim, dim * 3, bias=False) self.out = nn.Linear(dim, dim) self.edge_proj = nn.Linear(edge_dim, heads) if edge_dim else None self.drop = nn.Dropout(dropout) def forward(self, x, edge_index, edge_attr=None): N, H, Dh = x.size(0), self.h, self.dh q, k, v = self.qkv(x).view(N, 3, H, Dh).unbind(1) src, dst = edge_index[0], edge_index[1] score = (q[src] * k[dst]).sum(-1) / math.sqrt(Dh) # [M, H] if self.edge_proj is not None and edge_attr is not None: score = score + self.edge_proj(edge_attr) # 分段求 max（按 dst 归约） neg = torch.finfo(score.dtype).min mx = torch.full((N, H), neg, device=score.device, dtype=score.dtype) mx = mx.scatter_reduce(0, dst.unsqueeze(-1).expand(-1, H), score, reduce="amax", include_self=True) mx = torch.where(mx <= neg, torch.zeros_like(mx), mx) # 孤立节点保护 e = torch.exp(score - mx[dst]) den = torch.zeros(N, H, device=score.device, dtype=score.dtype) den = den.index_add_(0, dst, e) attn = self.drop(e / den[dst].clamp_min(1e-9)) out = torch.zeros(N, H, Dh, device=x.device, dtype=x.dtype) out = out.index_add_(0, dst, attn.unsqueeze(-1) * v[src]) return self.out(out.reshape(N, H * Dh))

三个坑必须点出来。第一个是clamp_min，孤立节点（入度为 0）的分母是 0，不加保护会出 NaN，而且这种 NaN 会顺着反向传播污染整个模型，排查起来极其痛苦。第二个是mx的孤立节点保护，长度为 0 的段求 max 得到的是负无穷，减出来是 NaN，同样要处理。第三个是index_add_在 CPU 上慢得离谱，在上面做小规模验证可以；真要训练，要么放到 GPU 上用它的原子加实现，要么换成专门的稀疏算子。

还有一个数值细节：e / den这种写法，如果某条边对应的差值是 -inf（异常输入），exp 后是 0，分母可能仍然正常，不会崩但会静默丢边。建议在开发阶段加一个检查，统计每层 softmax 后注意力权重为 0 的边比例，超过阈值就报警。

5. 训练流程、参数实验与效果对比

5.1 位置编码不能省

Exphormer 本身不提供位置编码，它假设外层框架已经处理了。这一点很重要，我见过有人直接把它当成一个完整的图 Transformer 用，结果所有节点的初始特征完全相同，模型退化成对结构不敏感的均匀聚合。

实际做法是在输入层加位置编码，常用的两种：拉普拉斯特征向量编码（取前几个非平凡特征向量，符号做随机翻转以避免歧义），以及随机游走返回概率编码（统计不同步长下回到自身的概率，作为结构指纹）。图规模大、特征分解代价高的时候，随机游走那种更实用，因为它只需要稀疏矩阵乘法，成本线性。

5.2 小批量与子图重采样

大图跑不动完整的图，就要采样子图训练。这里有个非常隐蔽的坑：子图上的扩展图必须重新采样。如果你在超图上采好扩展图，然后按节点子集去筛边，筛出来的子图里扩展边会大量缺失，连通性被破坏，长程信息彻底断掉，模型表现会莫名其妙地退化到接近只用局部边的水平。

正确做法是每采一个子图，在这个子图的节点集上重新跑一遍build_expander。这会带来额外的计算，但相比采样误差带来的损失，这点开销微不足道。我踩这个坑的时候调了三天超参都没找到原因，最后打印子图连通分量数量才发现问题。

5.3 参数组合实测记录

我用一个中等规模的数据集做了组对比，指标是验证集上的相对值（以基线为 100，越高越好），显存按单层前向峰值估算。数据量有限，只能看趋势，绝对值不要当真。

结论很清楚：d 从 4 提到 8 收益明显，8 提到 16 基本饱和；虚拟节点 k = 1 是最优性价比，k = 3 反而掉了；层数从 12 加到 16 有微小提升，继续加就过平滑了。这和张量化的混合时间估算对得上——d = 8 大约需要 22 跳，但虚拟节点把长程通道压缩成 2 跳，实际有效感受野远大于 16 跳。

5.4 大图上的显存观测

我把同一份配置放到三万节点的图上跑了一遍，稠密版本在 8 层、8 头时单卡直接 OOM，稀疏版本（d = 8，k = 1）单 card 跑到 40 层才顶到上限。这个倍率和我前面算的理论值有出入，原因是小图的显存被固定开销（词表、位置编码缓存）主导，负载没有跑满，图越大越接近理论比值。

6. 常见问题与排查技巧实录

6.1 问题速查表

6.2 我踩过的几个具体坑

第一个是边去重。原始图和扩展图、虚拟边的拼接可能产生重复边，重复边在分段 softmax 里会被重复计算，相当于给那条连接加了权重。小图影响不大，大图上重复率高的时候会明显改变注意力分布。我的处理是在拼接后做一次去重，代价很小。

第二个是种子管理。扩展图的随机性是我调试时的最大干扰源。一开始我图省事，每一层都用全局随机状态生成，结果两次运行结果差 3 个点，根本没法判断改动是否有效。后来我把 seed 分成三组：数据采样 seed、扩展图 seed、模型初始化 seed，分别控制，固定其中两组只调一组，实验才可复现。

第三个是梯度检查点的副作用。为了省显存我开了梯度检查点，结果发现训练时间翻了 2.5 倍，而且分段归约的部分在重算时会有额外的临时内存，省下来的显存没那么多。大图上该开还是要开，但如果你的图节点数在五千以下，别急着开，先用更小的扩展图度数试试。

第四个是虚拟节点和 batch 维度的顺序。批量训练时虚拟节点是每个样本各自独立的，不是跨样本共享的。我一开始把虚拟节点放在 batch 维度外面，相当于所有样本共用一个全局记忆槽，信息串了，指标虚高但泛化崩了——这个 bug 非常隐蔽，因为训练集上表现很好，只有验证集才暴露。

6.3 一个可以立刻试的小技巧

如果你现有的图 Transformer 已经跑得动，想低成本验证一下稀疏注意力值不值得投入，可以先不动模型，只在原有全连接注意力的基础上加一张扩展图和虚拟节点，把注意力掩码限定在“全连接保留，但只对这三类边计算”，用一个简单的布尔掩码跑一遍。这一步改动很小，如果指标不降或微升，说明长程依赖确实主要靠这几条通道承载，那再往完整的稀疏实现迁移就是稳赚；如果指标明显下降，说明你的任务对稠密注意力的精细建模依赖很强，得在边特征或者局部通道上多下功夫，不要盲目上稀疏。

我个人在大图任务上的体会是，稀疏化的收益主要不在单点精度，而在“能不能把图跑起来”这条硬门槛上。能跑起来之后，再靠虚拟节点数量、扩展图度数、层数这个小三角去调，比死磕注意力机制本身性价比高得多。至于扩展图的随机性，我的建议是接受它，用固定 seed 管理起来，不要试图找一个“最优”的扩展图——理论保证是建立在随机性之上的，你一旦去优化它，那些界就不成立了。