★ FEATURED ARTICLE 图自监督预训练这两年有点像个炼丹现场——手里捏着一堆没有标签的图数据位置信息、结构信息、节点属性全都有唯独缺标注而对比学习那套正负样本构造、图增广策略、动量编码器的组合拳调起来实在让人头大增广方式换一下下游指标就抖三抖负样本池小一点模型直接学塌。GraphMAE 是我觉得这几年在图预训练里思路最干净的一个方案它把 NLP 和视觉里已经验证过的掩码自编码器Masked Autoencoder搬到了图上随机遮掉一部分节点的特征让模型靠着邻居的信息把被遮住的部分一个字一个字还原出来。整个过程不需要负样本、不需要精心设计的图增广、不需要动量编码器只用一个编码器加一个轻量解码器就能跑起来。这篇文章我会从设计动机、核心机制、参数取舍、完整实操代码一路聊到实际落地时踩过的坑适合已经会用 PyTorch 和 PyG 跑通基础 GNN、想在图数据上做预训练或者表征学习的同学参考也适合正在纠结到底该用对比学习还是生成式预训练的从业者做技术选型判断。1. 为什么图预训练需要掩码重构这条路1.1 图神经网络的标签饥渴与对比学习的隐痛做图相关任务的同行大概都有同一个感受图结构好拿标签极贵。一张社交网络图、一份分子结构库、一张引文网络节点动辄几十万上百万但真正带标注的往往只有几百到几千个。更麻烦的是图上的标注常常需要领域知识——判断一个分子节点的毒性、判断一篇论文的领域分类都不是随便拉个人点两下就能标出来的。所以自监督预训练在图领域一直是刚需大家希望先在无标签图上把模型喂饱再拿少量标签做下游微调。前几年主流方案基本是图对比学习这条路用两个增广视图做正样本对再用一堆负样本把它们推开。方法本身没问题实际用起来有三个绕不开的痛点。第一是增广策略高度依赖数据节点丢弃、边扰动、属性掩码这几种增广在不同图上的效果差异极大同配图上表现好的配置换到异配图上可能直接掉点等于每换一个数据集就得重新调一轮增广组合工作量隐形但巨大。第二是负样本的依赖无论是对比池还是 in-batch 负样本都需要较大的 batch size 才能保证负样本足够难这对显存是硬约束图数据本来就吃显存再叠一个大 batch 很容易爆。第三是目标与下游任务的不一致对比学习优化的是实例判别学到的是区分性表征但下游节点分类要的是判别性加语义性的表征两者并不完全对齐。另一条更老的路是基于图自编码器GAE/VGAE的重构邻接矩阵。这条路的问题在于任务太简单。给定一个节点的邻居集合让它还原邻接矩阵模型很容易学到一个接近度或者局部连通性的平凡解——因为邻接矩阵里有大量零元素重构损失被这些零主导模型只要把谁和谁不相连学好就能把 loss 压下去真正有用的高阶语义信息反而没学到。我试过在引文网络上用 GAE 做预训练下游节点分类的提升微乎其微有时候甚至不如直接训练 GCN 加早停。1.2 GraphMAE 的整体设计取舍GraphMAE 的出发点很直接既然邻接重构太容易、对比学习太依赖设计那就换一个难度适中且语义信号强的重构目标——节点特征。具体做法是随机选一部分节点把它们的特征向量整体替换成一个可学习的掩码标记然后让编码器在整张图上做消息传递最后由解码器根据编码后的隐表示把被掩码节点的原始特征还原出来。损失只在被掩码的节点上计算。这个设计里有两个关键取舍值得展开说。第一个是编码器看整张图而不是只看未掩码部分。视觉里的 MAE 之所以只编码可见 patch是因为图像冗余极高一张图里 25% 的 patch 已经足够推断出全局内容而且只编码可见部分能把计算量降下来。但图不一样图节点的特征冗余度低得多如果只编码可见节点被掩码的节点就完全失去了参与消息传递的机会编码器学到的表示会天然偏向可见子图。GraphMAE 选择让编码器处理包含掩码标记在内的全部节点掩码标记本身也作为一个特殊的占位特征参与聚合这样每个节点的表示都融合了邻居的上下文解码器拿到的隐表示质量更高。第二个是重构目标是特征不是结构。这个选择背后有个很实在的观察图节点特征本身往往携带了很强的语义信号比如引文网络里节点的词袋特征、分子图里的原子类型和属性这些特征维度之间是有关联的模型想还原它们就必须理解节点在整个图中的角色和邻域语义。而且特征重构的难度天然可控——掩码比例一调任务难度就变了不像邻接重构那样存在大量平凡零元素把损失稀释掉。1.3 与后续变体的边界分清楚需要说明的是GraphMAE 本身是一个相对精简的框架后面还有 GraphMAE2 这类增强版本引入了去噪解码、随机重掩码、多视图隐空间一致性等机制。我个人的判断边界是这样的如果你手上的图是中小规模几万到几十万节点、同配性较强、特征维度在几百到几千之间那么原版 GraphMAE 的性价比最高代码量小、超参少、训练稳定如果你追求极致的下游指标或者图上噪声边特别多、异配性明显那可以考虑引入去噪解码那一套代价是代码复杂度和调参难度都会上去。这篇文章主要围绕原版框架展开因为它才是理解整条技术脉络的起点。2. 核心机制拆解掩码、编码、解码、重构损失2.1 掩码比例与掩码标记的填充方式掩码策略看起来是最简单的一环实际上坑不少。GraphMAE 用的是随机节点掩码也就是从 N 个节点里无放回地抽一部分把这些节点的特征整体替换掉。注意这里是整体替换而不是逐维度掩码这一点和 BERT 的 token 掩码逻辑更像而不是像图像 MAE 那样按 patch 遮挡一块区域。掩码比例这个参数常见取值在 0.5 到 0.75 之间比 BERT 常用的 15% 高出一大截。原因在于图的节点特征冗余度远低于文本 token文本里一句话的每个词都受上下文强约束遮 15% 就已经很有挑战了而图上一个节点的特征向量如果只遮 15% 的节点剩下的 85% 节点信息非常充分模型几乎不需要真正理解邻域语义就能把特征猜个八九不离十任务就变得太容易了学到的表示也就没什么价值。我实测下来在同配性较强的引文网络上0.5 是个比较稳的起点分子图这类节点数较少的图建议往 0.3 到 0.5 靠因为图本身就小遮太多会导致连通性被破坏。填充方式也是一个容易被忽略的细节。常见做法有三种填零、填一个可学习的向量、把可学习向量加到原特征上。填零的问题在于它会引入一个明显的分布偏移——零向量在特征空间里是一个特殊点模型很容易学会只要看到零就知道这里被遮了然后偷懒用一个固定的输出去重构它等于绕过了任务。填可学习向量的效果最好因为掩码标记会随着训练自适应成一个信息真空的表示编码器会自然地学会从邻居聚合信息来填补它。代码里我通常这么写x_masked x.clone() x_masked[mask_nodes] self.enc_mask_token # nn.Parameter, shape [in_dim]有个小技巧是把这个参数初始化成全零或者很小的随机值训练初期不要让它太强否则会主导整张图的特征分布。2.2 编码器与解码器的不对称设计GraphMAE 的编码器和解码器是刻意做成不对称的。编码器一般用两层 GNNGCN、GAT、GIN、GraphSAGE 都行隐藏维度 256 到 512每层后接 LayerNorm 和 ReLU再加 dropout解码器则轻得多通常是一层 GNN 加一个线性映射或者干脆用 MLP 加一层图卷积。为什么要让解码器这么弱这其实是一个防止信息泄漏的设计。如果解码器足够强它完全可以绕过编码器学到的语义表示直接从掩码标记的位置信息和图结构里硬凑出答案那编码器就学不到东西了。把解码器压薄等于逼着编码器在隐空间里把语义信息组织好解码器只负责把这个隐表示翻译回特征空间。这个思路和 MAE 里解码器比编码器浅是一个道理。这里还有个细节叫re-mask重掩码在编码器输出隐表示之后、送进解码器之前再对隐表示做一次随机掩码用另一个可学习的标记替换。这个操作的直觉是进一步增加任务的难度防止解码器直接看到被掩码节点自己的隐表示然后抄答案。原版里 re-mask 的比例是个超参我一般设成和主掩码比例相同或者略高一些0.5 到 1.0 之间在小图上效果比较明显大图上收益会递减。还有一点必须提醒归一化层用 LayerNorm不要用 BatchNorm。预训练阶段我们通常是一整张图或者一个大子图作为一个 batchBatchNorm 的统计量在预训练和下游微调之间会严重不一致尤其是微调阶段 batch 很小的时候BN 的 running stats 会直接把性能拖垮。这个坑我在早期项目里踩过当时预训练 loss 降得很漂亮下游一微调指标反而不如随机初始化排查了两天才定位到是 BN 的问题。2.3 重构特征而非结构的取舍逻辑前面提过邻接重构的平凡解问题这里再往深一层说。重构邻接矩阵本质上是让模型预测两个节点之间有没有边这个任务的监督信号是二值的、极度稀疏的模型很容易学到一个基于节点度的启发式。而重构节点特征监督信号是高维稠密的模型必须真正理解这个节点在当前邻域里应该长什么样才能把几百维的特征向量还原到余弦相似度较高的水平。反过来讲特征重构也有它的局限如果图上的节点特征本身信息量极低比如全是 one-hot 的 ID或者维度只有个位数那特征重构能提供的信号就很有限这时候 GraphMAE 的优势会打折扣。我在一个只有 8 维类别特征的图上试过预训练几乎没有带来增益后来改成先做特征工程、把节点属性扩展成多维统计量效果才出来。所以选型的时候先看一眼特征维度低于 32 维的图要慎重。2.4 缩放余弦误差为什么不直接用 MSE损失函数是 GraphMAE 里最容易被低估的一环。原版用的是缩放余弦误差Scaled Cosine Error公式大致是这样的L (1/N_m) * Σ_{i∈masked} ( (1 - cos(z_i, x_i)) / τ )^γ其中 z_i 是解码器输出x_i 是原始特征τ 是缩放因子默认 1γ 是幂次默认 2N_m 是被掩码节点数。为什么不用 MSE关键在于特征尺度的异质性。真实图数据的特征维度之间数值范围差异极大比如某一维是计数0 到几千另一维是二值标记0 或 1用 MSE 的话大数值维度会完全主导损失模型只顾着把那几维还原好小数值维度基本被忽略。余弦相似度只看方向不看模长天然对这种尺度差异免疫每个节点的特征被归一化到单位球面上各维度的重要性更均衡。τ 和 γ 这两个参数的作用也不难理解。γ 大于 1 的时候损失函数对小误差的惩罚被压低、对大误差的惩罚被放大相当于让模型更关注那些难还原的节点避免它在容易的样本上刷分。τ 则控制整体损失的尺度影响等效学习率一般不用动保持 1 就行。我在几个数据集上试过 γ 从 1 到 3 的变化2 是比较通用的选择γ 太大比如 4会导致训练前期 loss 剧烈震荡梯度不太稳。还有一个前置条件原始特征最好先做行归一化。因为余弦损失只关心方向如果不同节点的特征模长差异巨大归一化之后的信息会更干净。这一步在代码里就一行x F.normalize(x, p2, dim1)但它对结果的影响不小实测下来能带来一到两个点的下游提升。3. 实操落地从零跑通一个 GraphMAE 预训练流程3.1 环境与数据准备环境上没什么特殊的PyTorch 加 PyTorch Geometric 就够版本上建议 torch 2.x 配 PyG 2.3 以上主要是为了避免老版本里GCNConv的稀疏矩阵实现对 batch 维度的兼容问题。pip install torch torchvision pip install torch_geometric pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.1.0cu118.html数据准备阶段有三件事要做顺序不能乱。第一是读图并转成 PyG 的 Data 对象得到x节点特征矩阵形状[N, d]和edge_index形状[2, E]。第二是特征清洗把全零行、NaN、inf 处理掉全零行在余弦损失下会直接导致normalize之后得到零向量反向传播时梯度异常必须清掉或者用邻域均值填。第三是特征归一化做行 L2 归一化。import torch import torch.nn.functional as F from torch_geometric.datasets import Planetoid dataset Planetoid(root./data, nameCora) data dataset[0] x data.x.float() x torch.nan_to_num(x, nan0.0, posinf0.0, neginf0.0) zero_rows (x.abs().sum(dim1) 0) if zero_rows.any(): # 用邻居特征均值填充全零行 row, col data.edge_index for v in zero_rows.nonzero(as_tupleTrue)[0]: nbrs col[row v] if len(nbrs) 0: x[v] x[nbrs].mean(dim0) x F.normalize(x, p2, dim1) data.x x这段清洗代码看着平淡但在真实数据上非常必要。我在一个真实业务图上就遇到过超过 3% 的节点特征是全零预训练 loss 一直卡在某个值下不去排查后发现是normalize对零向量返回了零梯度这些节点相当于完全死在了训练里。3.2 模型配置参数逐项拆解先把核心超参列成表方便对照着调参数常用取值作用调整建议掩码比例 mask_ratio0.3 ~ 0.75控制任务难度大图偏 0.5小图/分子图偏 0.3隐藏维度 hid_dim256 / 512表示容量与特征维度同量级或 2 倍编码器层数2 ~ 3感受野同配图 2 层够异配图可试 3解码器层数1防止绕过编码器不建议超过 2re-mask 比例0.5 ~ 1.0增加难度小图开大图可关γ幂次1 ~ 3难样本权重默认 2震荡就降到 1dropout0.2 ~ 0.3正则小图调高大图调低学习率1e-3 ~ 5e-3优化步长配合 warmupweight decay0 ~ 1e-5正则基本不用或设极小隐藏维度的选择有个经验公式取节点特征维度 d 的 1 到 2 倍并对齐到 64 的整数倍方便 GPU 计算。比如 Cora 的 d 是 1433我一般直接取 512再大显存就吃不消了再小表示容量不足。编码器的选择要看图的同配性。如果是引文网络、社交网络这类同配图相邻节点标签相近GCN 就够用简单稳定如果是分子图、异配性强的图GAT 或 GIN 更合适GAT 的注意力机制能自动降低噪声邻居的权重如果是大规模图GraphSAGE 的邻居采样版本更实用。3.3 完整训练循环实现把模型定义、掩码逻辑、损失函数串起来下面是核心代码。import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import GCNConv, GATConv class GNNEncoder(nn.Module): def __init__(self, in_dim, hid_dim, num_layers2, conv_typegcn, dropout0.2): super().__init__() self.convs nn.ModuleList() self.norms nn.ModuleList() for i in range(num_layers): d_in in_dim if i 0 else hid_dim if conv_type gcn: self.convs.append(GCNConv(d_in, hid_dim)) elif conv_type gat: self.convs.append(GATConv(d_in, hid_dim // 4, heads4, concatTrue)) self.norms.append(nn.LayerNorm(hid_dim)) self.dropout dropout def forward(self, x, edge_index): for conv, norm in zip(self.convs, self.norms): x conv(x, edge_index) x norm(x) x F.relu(x) x F.dropout(x, pself.dropout, trainingself.training) return x class GraphMAE(nn.Module): def __init__(self, in_dim, hid_dim512, enc_layers2, dec_layers1, mask_ratio0.5, remask_ratio0.5, conv_typegcn, dropout0.2): super().__init__() self.mask_ratio mask_ratio self.remask_ratio remask_ratio self.enc_mask_token nn.Parameter(torch.zeros(in_dim)) self.dec_mask_token nn.Parameter(torch.zeros(hid_dim)) self.encoder GNNEncoder(in_dim, hid_dim, enc_layers, conv_type, dropout) self.decoder GNNEncoder(hid_dim, in_dim, dec_layers, conv_type, dropout) def random_mask(self, n, ratio, device): num int(n * ratio) perm torch.randperm(n, devicedevice) mask torch.zeros(n, dtypetorch.bool, devicedevice) mask[perm[:num]] True return mask def forward(self, x, edge_index): n x.size(0) mask self.random_mask(n, self.mask_ratio, x.device) x_in x.clone() x_in[mask] self.enc_mask_token h self.encoder(x_in, edge_index) if self.remask_ratio 0: remask self.random_mask(n, self.remask_ratio, x.device) h h.clone() h[remask] self.dec_mask_token z self.decoder(h, edge_index) return z, x, mask def scaled_cosine_error(z, target, tau1.0, gamma2.0, eps1e-8): z F.normalize(z, p2, dim-1, epseps) t F.normalize(target, p2, dim-1, epseps) cos (z * t).sum(dim-1) return ((1.0 - cos) / tau).clamp(min0).pow(gamma).mean()训练循环注意三点一是每个 epoch 都重新采样掩码位置这一点和 dropout 类似属于随机正则不要固定一套掩码二是损失只在被掩码的节点上算未掩码节点的输出直接丢进垃圾桶不要图省事对全部节点求损失那样会稀释信号三是加一个学习率 warmup前 5% 到 10% 的步数线性升温能让训练早期更稳。model GraphMAE(in_dimdata.num_features, hid_dim512, mask_ratio0.5) optimizer torch.optim.Adam(model.parameters(), lr1e-3, weight_decay0.0) epochs 500 warmup max(1, int(epochs * 0.05)) for epoch in range(epochs): model.train() lr 1e-3 * min(1.0, (epoch 1) / warmup) for g in optimizer.param_groups: g[lr] lr optimizer.zero_grad() z, x_orig, mask model(data.x, data.edge_index) loss scaled_cosine_error(z[mask], x_orig[mask], tau1.0, gamma2.0) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) optimizer.step() if epoch % 50 0: print(fepoch {epoch}, loss {loss.item():.4f}, lr {lr:.2e})我在 Cora 上跑这套配置loss 大致会从 0.35 左右缓降到 0.10 到 0.15 区间。如果 50 个 epoch 后 loss 还在一开始的值附近不动先检查特征归一化是不是漏了或者学习率是不是太小。3.4 下游线性探针与微调评估预训练完之后怎么评估这决定了你能不能判断预训练到底有没有用。最省事的做法是线性探针冻结编码器参数只在编码器输出的隐表示上接一个单层线性分类器用少量标签训练。这个协议的好处是能干净地衡量表示质量——如果隐表示本身就把类别分开了一个线性层就能达到不错的准确率。class LinearProbe(nn.Module): def __init__(self, hid_dim, num_classes): super().__init__() self.fc nn.Linear(hid_dim, num_classes) def forward(self, h): return self.fc(h) # 冻结编码器 for p in model.encoder.parameters(): p.requires_grad False model.encoder.eval() with torch.no_grad(): h model.encoder(data.x, data.edge_index) probe LinearProbe(h.size(1), dataset.num_classes).to(h.device) opt torch.optim.Adam(probe.parameters(), lr1e-2, weight_decay5e-4) for _ in range(200): probe.train() opt.zero_grad() out probe(h) loss F.cross_entropy(out[data.train_mask], data.y[data.train_mask]) loss.backward() opt.step() probe.eval() with torch.no_grad(): pred probe(h).argmax(dim-1) acc (pred[data.test_mask] data.y[data.test_mask]).float().mean() print(ftest acc: {acc.item():.4f})线性探针跑出来的数字一般比完整微调低几个点但它的稳定性好适合用来做不同预训练配置之间的横向对比。关键的一点是要设一个随机初始化的基线同样架构、同样训练轮数、不做预训练直接随机初始化编码器然后跑线性探针。只有预训练跑出来的数字明显高于这个基线才能说明预训练真的起了作用。我见过太多人拿着 80% 的准确率说预训练有效其实随机初始化也能到 78%提升全在噪声范围内。4. 参数调优与常见坑排查实录4.1 掩码比例与损失函数的经验取值掩码比例是影响最大的单一超参值得单独花时间调。我的做法是固定其他配置在[0.3, 0.5, 0.65, 0.75]这四档里各跑一遍用线性探针的准确率来选。经验规律是图的平均度越高、特征维度越大可以承受的掩码比例越高。因为高平均度意味着被掩码节点有更多邻居提供信息高特征维度意味着语义冗余更多两者都给重构任务提供了线索。反过来稀疏图平均度小于 3和小特征图要从 0.3 起步遮太多会让连通性断裂模型根本无从推断。γ 这个幂次参数容易被忽略但它对最终效果的影响有时能到一两个点。γ 等于 1 的时候损失函数是标准的余弦距离比较温和γ 等于 2 的时候误差被平方放大模型更聚焦于难样本通常效果更好γ 大于 3 就开始不稳了前期 loss 震荡剧烈。有一个细节γ 变大之后等效的梯度尺度也变大了这时候需要相应调小学习率否则容易发散。re-mask 的比例我一般在 0.5 到 0.8 之间取。设为 0关掉也能训但在小图节点数小于 5000上关掉之后线性探针通常掉 0.5 到 1.5 个点因为小图上信息密度高解码器太容易作弊。大图上这个技巧的收益就不明显了可以为了省事直接关掉。4.2 表征塌缩与训练不收敛的排查表下面这张表是我实际排查问题时总结的按现象查原因命中率挺高现象可能原因排查与解决loss 卡在某值不动特征未归一化或全零行过多检查x的行范数分布清洗全零行loss 前期震荡剧烈γ 过大或学习率过高γ 降到 1lr 降到 3e-4训练 loss 很低但下游很差过拟合预训练轮数过多减半 epochs加 dropout下游不如随机初始化归一化层用了 BatchNorm全部换成 LayerNorm隐表示各维方差极小表征塌缩检查解码器是否过深降低 re-mask部分节点输出恒定孤立节点无邻居给孤立节点加自环或单独处理线性探针梯度爆炸隐表示未归一化探针输入前加一层 LayerNorm表征塌缩这一项要重点说。虽然 GraphMAE 不像对比学习那样容易出现完全塌缩但如果解码器设计得过强比如三四层 GNN编码器学到的表示会趋向于各向同性不同节点的隐向量余弦相似度都接近 1下游完全没法用。判断方法很简单随机抽 1000 个节点算隐表示的余弦相似度矩阵如果非对角元素的均值超过 0.9基本就可以判定塌缩了。解决办法是砍掉解码器层数、提高掩码比例、或者适当加大 dropout。4.3 小图、异构图、大图的适配改造小图场景节点数少于 2000比如多数分子图数据集最麻烦的是批量化。单张图作为一个样本太小训练效率极低正确的做法是把多张小图打包成一个 batch用 PyG 的Batch做图级别拼接同时在每张小图内部独立做掩码。这里有个陷阱拼接之后mask的索引是全局的计算损失时必须把 mask 和原始特征对齐否则会串到别的图上。我一般会在forward里按batch向量分段处理每张图单独采样掩码。异配图场景相邻节点标签差异大比如某些推荐场景的二分图需要换编码器。GCN 的同配假设在异配图上直接失效此时用 GAT 让注意力自适应地选择邻居或者干脆用两层 MLP 加一跳邻居聚合的结构。另一个思路是在预训练阶段引入边掩码作为辅助任务让模型同时理解结构和特征但这个改动已经超出原版框架了。大图场景千万级节点的核心矛盾是显存。原版 GraphMAE 假设一次性编码整张图这个假设在大图上直接破产。可行的改造是用邻居采样每次只采一个子图做前向但采样时要注意保证被掩码节点的 k 跳邻居尽量被采到否则那个节点的重构任务就变成了无中生有信号是噪声。我一般设采样邻居数为[15, 10]两层batch 里放 8 到 16 个种子节点跑出来的效果比全图训练差一点但至少跑得动。5. 工程化落地与效果评估的几个现实问题5.1 评估协议怎么选才不骗自己评估协议的选择直接决定了你对预训练效果判断的准确性。线性探针最干净但数字偏低完整微调数字高但受很多因素干扰学习率、训练轮数、早停策略不同配置之间可比性差。我的建议是两个都跑线性探针用来做超参筛选因为它的方差小、对超参不敏感完整微调用来汇报最终数字因为它更接近真实业务场景。还有第三种协议是少样本评估只用 1 个、5 个、10 个标签每类来训练下游分类器。这个协议最能体现预训练的价值因为标签越少预训练学到的表示占比越大。我在几个项目里观察到全量标签下 GraphMAE 相比随机初始化的提升可能只有 1 到 2 个点但在每类 5 个标签的设置下提升能到 8 到 15 个点。如果你的业务场景本身就缺标签那这个指标才是真正该关注的。评估里还有一个必须统一的细节特征预处理必须和预训练时保持一致。预训练时做了 L2 行归一化下游评估时如果忘了做分布就变了隐表示的质量会明显下降。我在代码里习惯把归一化逻辑封成一个preprocess(x)函数预训练和下游共用避免遗漏。5.2 大规模图上的显存与速度控制显存控制除了采样还有几个实用的手段。第一是混合精度训练用torch.cuda.amp把前向和反向压到 fp16显存能省 30% 到 40%对收敛的影响很小注意 loss 计算里的normalize和.pow()要保持 fp32否则数值精度不够。第二是特征降维如果原始特征维度是几千维可以先做一次 PCA 或随机投影降到 256 维重构目标的维度下降之后解码器的输出层参数量和中间激活都大幅减少。实测下来降维到 256 对下游指标的影响通常在 1 个点以内但显存能省一半。第三个手段是梯度累积因为图预训练时 batch 通常不大全图或者少量子图梯度累积能让你在显存不变的前提下用更大的等效 batch这对训练的稳定性有好处。具体做法是每 N 步才optimizer.step()一次中间不断累加梯度。速度方面PyG 的GCNConv在中小图上已经足够快瓶颈往往在特征矩阵的拷贝上。注意x_in x.clone()这行如果每个 epoch 都 clone 一个大矩阵内存分配开销不容忽视。可以把 clone 换成预先分配好缓冲区的index_put_或者用一个固定的 mask 缓冲区反复写入实测能省 10% 到 20% 的时间。5.3 什么情况下不该用 GraphMAE最后聊聊边界这比什么时候用更重要。有三种情况我建议直接放弃 GraphMAE别浪费时间。第一种是节点特征信息量极低的图。比如特征全是 one-hot 的节点 ID或者维度小于 16 的粗粒度类别特征。这类图上特征重构任务几乎学不到语义模型只能记住 ID 和邻居的对应关系学到的表示在下游迁移性极差。这种情况应该先做特征工程或者转向基于结构而非特征的自监督方法。第二种是图规模巨大且对训练成本敏感的场景。如果图有几亿节点而你的算力预算只够跑几十个 epoch那 GraphMAE 的采样改造和调参成本可能比收益高。这种情况下不如直接用无监督的结构嵌入比如随机游走类的浅层方法配合简单 GNN性价比更高。第三种是下游任务与节点特征强相关但标签充足的场景。如果每个类别都有上万条标注预训练带来的边际收益会快速衰减直接监督训练加合适的正则可能就够了。预训练真正的价值区间是图大、标签少、特征有意义这个三角。我在实际项目里的体会是GraphMAE 最舒服的使用姿势是先花半天时间把数据和特征清洗干净用默认配置掩码 0.5、γ 取 2、两层 GCN、隐藏 512跑一版预训练然后用线性探针和随机初始化基线对比如果提升超过 3 个点就继续调掩码比例和 re-mask如果提升在 1 个点以内先别急着调参回去检查特征工程和归一化流程八成是数据侧的锅。这套流程我在几个不同类型的图上跑下来成功率比一上来就猛调超参高得多。另外分享一个小技巧预训练结束时把编码器的输出隐表示保存下来下游无论做线性探针、聚类、还是当特征喂给树模型都能直接复用省得每次都重新前向一遍。
阅读完成 · 觉得有帮助？
