ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

URFormer:关系词表与专家路由实现异构图零样本迁移

URFormer:关系词表与专家路由实现异构图零样本迁移

做异构图模型的同学大概都有过这种体验:业务方甩过来一张新图,schema 里躺着二十几种节点类型、上百种边语义,你手里的模型只认得训练时见过的那几种关系,于是又是一轮重写特征、重训 embedding、重调参。一个模型,任意关系类型这句话之所以让人眼前一亮,就是因为它戳中了这个反复返工的痛点。我这里要拆的这套框架,内部代号叫 URFormer(统一关系建模框架),核心主张是:把多张异构信息网络并成一张全局异构图,用关系类型词表把每条元关系编码成可组合的向量,再用一组按关系类型路由的轻量专家来承载不同关系的语义差异。模型预训练一次,下游遇到没见过的关系类型时不需要重新训练主干,也能直接干活。这篇文章适合三类人:手上异构关系超过二十种、已经被 schema 演进折磨过的算法工程;想搞清楚"零样本迁移到新关系类型"具体怎么落地、而不是只看个指标的读者;以及准备自己复现一版做对比实验的研究者。下面我会把设计取舍、机制细节、可抄的训练配置和踩过的坑一次讲完。

1. 为什么"一个模型,任意关系类型"这句话值得较真

先把问题摆到台面上。现实里的图数据几乎不存在"节点只有一种、边只有一种"的情况。电商场景里,用户、商品、店铺、品类、品牌、优惠券互相连接,边有浏览、加购、下单、退款、收藏;风控场景里,设备、账号、手机号、银行卡、收货地址交织,边有登录、绑定、转账;内容场景里,作者、文章、话题、评论互相引用。这些图本质上都是异构信息网络(Heterogeneous Information Network,HIN),它们的骨架由节点类型和关系类型共同定义,而不是由节点本身定义。

过去几年主流做法是"一场景一模型":给每种关系单独设一套 embedding 或者一条独立的注意力通道,训练时把见过的关系类型都塞进去。短期指标确实好看,长期维护成本高得离谱。最直接的后果是模型参数量随关系类型数量线性增长,而真正高频的关系只有那么七八种,剩下几十种长尾关系每种可能只有几千条边,单独学一套参数根本学不动。更麻烦的是业务侧不会等你,新品类、新互动行为、新设备指纹每个月都在加,模型每次都要重训一遍。

1.1 先把"关系类型"这个概念掰开

很多人嘴上说关系类型,实际上心里想的是边的名字。这两者不是一回事。在这套框架里,一条关系类型被定义成一个三元组:源节点类型、边语义、目标节点类型,学术上常叫元关系(meta-relation)。举个具体的例子,(用户, 购买, 商品)和(用户, 点击, 商品)是两条不同的关系类型;(用户, 购买, 商品)和(店铺, 售出, 商品)也是两条不同的关系类型,尽管边语义有交集。

这个定义看着朴素,但它直接决定了问题的规模。假设一张图里有 |A| = 10 种节点类型,|R| = 12 种边语义,那么理论上可能出现的关系类型数量是 10 × 12 × 10 = 1200 种,而实际在数据里出现的可能只有一百多种。也就是说,关系类型的空间是组合爆炸的,但真实观测到的只是其中稀疏的一小片。这个"空间大、观测稀"的结构,恰恰是"任意关系类型"能靠组合泛化做到的物理基础——如果每个关系类型都是完全独立的新事物,那零样本迁移根本无从谈起。

我把这个观察总结成一句话:关系类型不是原子,是可以拆的。拆成源类型、边语义、目标类型三段之后,没见过的组合就能由见过的零件拼出来。整套方案的设计几乎都是围绕这句话展开的,后面讲 tokenizer 的时候你会看到它是怎么落到实现上的。

1.2 关系类型爆炸带来的三个真实麻烦

第一个麻烦是长尾关系的表征学不动。单独给一条关系类型配一套参数,样本量少的那些关系就永远停留在随机初始化附近。我实测过一个场景:最高频的关系有 800 万条边,最冷门的只有 300 条边。如果用独立参数,冷门关系的召回率基本等于没训。而这类冷门关系往往业务价值不低——比如"退款"关系的样本天然就比"下单"少,但它对风控判断的价值更高。

第二个麻烦是schema 演进导致模型失效。业务新增一种互动行为,比如"收藏到合集",这条关系的节点类型组合是(用户, 收藏到, 合集),训练集里从来没有过。旧模型的常规处理是:把它当成未知边类型丢掉,或者退回同构 GNN 用节点特征硬扛。丢掉信息是浪费,硬扛则意味着模型对这个信号完全无感。这两种处理在实际业务里都很难交代。

第三个麻烦是跨域迁移时语义对不上。A 业务里是(买家, 下单, 货品),B 业务里是(用户, 购买, 商品),人类一眼就知道是一回事,但模型看到的是两个不同的符号 ID,无法共享任何知识。这意味着每开一个新业务线,就等于从零开始攒数据、训模型。这套框架想解决的,本质上是让模型学到"关系语义"这件事本身,而不是死记硬背关系类型的编号。

1.3 这套思路适合谁、不适合谁

说清楚适用边界比吹效果更重要。适合的场景有四个特征:关系类型数量多(超过 20 种)、长尾关系占比高、schema 会持续演进、存在跨域迁移需求。推荐系统、风控图谱、内容理解、学术图谱这类场景基本都符合。

不适合的场景也很明确。如果你只有一张同构图,节点一种、边一种,那这套东西的复杂度纯属浪费,直接上常规 GNN 更省事。如果你的图极其巨大但只跑少数几种固定关系,且对在线延迟卡到毫秒级,那么全局异构图带来的采样开销和路由开销可能不划算。还有一种情况要谨慎:如果你的关系类型之间语义差异极大,且每种关系的样本都很充足,那"共享专家"带来的收益会被稀释,独立建模反而更直接。我个人的判断标准是——先看长尾关系的边数占比,如果低于 5%,收益会很有限;如果在 15% 以上,值得试。

2. 架构选型:为什么是"全局异构图 + 关系词表 + 专家混合"

这套方案由三个部件咬合而成:把所有任务图并成一张全局异构图、用一个关系类型 tokenizer 把元关系编码成可组合向量、用一组按关系路由的轻量专家承载语义差异。这三件事单独拿出来都不新鲜,新鲜的是它们的组合方式。我在下面逐个说清楚为什么这么选,以及为什么不选另外几条看起来更直接的路。

2.1 把多张图并成一张全局异构图的三个好处

最直觉的做法是每张子图单独训一个模型,最后做集成。这个路子在小规模下能跑,但它有三个致命问题:参数量翻倍、跨图知识无法共享、推理时要维护一堆模型。合并成一张全局异构图之后,这三个问题一次性解决。

第一个好处是节点表征可以跨场景共享。同一个用户在推荐图里是买家,在内容图里是读者,在风控图里是账号主体。如果分开建模,这三份表征互不相通;合并之后,它们是同一个节点的一个 token,只是通过不同的关系类型连出去。这一点在冷启动场景下价值明显——新用户在主站没有行为,但在另一个场景里有历史,合并图能把这个信号带过来。

第二个好处是关系类型词表只需要维护一份。所有场景的元关系都进同一张词表,词表规模可控(实测一百多个元关系就够覆盖大部分业务),而且新增关系只需要往词表里加几个组合,不需要改模型结构。

第三个好处是采样可以在一个 batch 里混着来。训练时从一个全局图里按元关系分层采样,一个 batch 里可能同时有推荐样本、内容样本、风控样本。这种混合采样自带正则效果,我实测下来比单场景训练更不容易过拟合,尤其是在小场景上。

代价当然也有。全局图的邻居采样逻辑要写得足够稳,否则一个场景的超级节点会把 batch 撑爆;另外全局图的边表规模通常比单场景大一到两个数量级,对存储和加载有要求。我的处理办法是:节点特征存成按类型分片的稀疏矩阵,边表按元关系分文件,加载时按需 mmap,不要一次性全塞进内存。

2.2 关系类型 tokenizer:把一条关系拆成三段可学习的向量

这是整套方案里我认为最关键的一个设计。它的做法很朴素:一条元关系 (src_type, edge_type, dst_type) 不去学一个整体 ID 的 embedding,而是拆成三段分别查表,再拼起来过一个小的投影层。

为什么这么做?因为整体 ID 的 embedding 在面对新关系类型时只能随机初始化,而拆成三段之后,新组合的每一段大概率都在训练集里出现过。举例:(用户, 收藏到, 合集)这条关系在训练集里完全没出现过,但"用户"作为源类型出现过无数次,"收藏到"这个边语义可能在别的组合里出现过(比如(用户, 收藏到, 商品)),"合集"作为目标类型也可能出现过(比如(文章, 属于, 合集))。三段都有已知信息,拼出来的向量就不是随机的,模型能做出合理推断。

具体实现上,三段各有一个 embedding 表,维度设为 d/3 或者各 d/2 再加投影都行,我用的是各 128 维、拼接后过一层 256 维的线性投影加 LayerNorm。投影层的存在很必要,因为三段拼接后的分布和单段 embedding 差别很大,不归一化的话后面路由会抖。另外我建议给每一段额外乘一个可学习的重要性权重(三个标量,softmax 归一化),这样模型能自己调节"边语义更重要还是节点类型更重要"。在我的实验里,边语义的权重普遍偏高,符合直觉。

注意:三段 embedding 的表要分开建,不要把 src_type 和 dst_type 塞进同一张表。虽然它们都是节点类型,但扮演的角色不同,共享表会让"用户作为源"和"用户作为目标"混成一个向量,实测掉点 1 到 2 个百分点的 NDCG。

2.3 MoT 专家:不是给每种关系训一个模型,而是让关系挑专家

MoT 是 Mixture of Thought Experts 的缩写,本质上是一种稀疏专家混合。它的结构是:主干 Transformer 的所有参数共享,只在 FFN 层注入若干组低秩适配器(LoRA 或 adapter),每组称为一个"专家";前向时,根据当前 token 所属的关系类型向量,路由网络选 top-k 个专家,加权求和输出。

这个设计解决了"差异要建模、参数又不能爆炸"的矛盾。关系类型之间确实存在语义差异——"浏览"和"下单"对用户兴趣的表达强度完全不同,用同一套 FFN 处理会互相干扰。但差异又不是完全正交的,"点击"和"收藏"都是弱反馈,共享一套参数完全合理。让路由网络自己学哪些关系该共享、哪些该分开,比人工划分关系簇靠谱得多。

为什么用 LoRA 而不是给每个专家配完整的 FFN?算笔账就清楚了。假设隐藏维度 256、FFN 中间层 1024,一个完整 FFN 的参数量是 256×1024×2 ≈ 52 万。8 个专家就是 420 万,乘上 4 层是 1680 万,还只算了 FFN。改成 LoRA(rank=8),单个专家单层的参数量是 8×(256+1024) ≈ 1 万,8 个专家 4 层总共才 33 万左右,差了 50 倍。而且 LoRA 的初始化为零矩阵,训练初期等价于没有专家,不会破坏主干已经学好的表征,这对增量上线新关系类型特别友好。

路由网络的设计也有讲究。输入是关系类型向量(tokenizer 的输出)加一个可学习的全局偏置,输出是专家数量维的 logits,取 top-2 后做 softmax 加权。我试过 top-1,训练不稳定,因为路由一旦早期选错就再难纠正;top-3 以上收益不明显但显存上去了。top-2 是性价比最高的点。另外必须加负载均衡损失,否则路由会坍缩到少数几个专家上,这个坑我在第 5 节会细讲。

2.4 四条路线的横向对比

为了让你直观感受取舍,我把常见的四种做法摆在一起比一下。这张表是我自己在选型时整理的,数值来自同一个小规模基准(约 120 万元关系边、26 种元关系、留出 5 种元关系做零样本测试)。

方案参数随关系类型增长未见关系类型的支持方式零样本 NDCG@10显存占用(相对值)主要问题
独立 embedding 的 GNN 基线线性增长无,直接丢弃0.211.0长尾关系学不动
分层注意力(类型级 + 节点级)近似线性退回类型无关注意力0.291.4schema 一变就退化
单塔 Transformer,关系只做偏置常数有偏置但无专家区分0.331.1语义差异大的关系互相干扰
全局图 + 关系词表 + MoT常数组合泛化 + 路由推断0.411.6采样与路由调参成本高

看这张表要注意,最后一行领先的 0.08 不是在所有场景都能复现的。它的优势集中在两类样本上:一是训练时完全没出现过的元关系,二是长尾元关系(边数少于 5000 的那些)。在高频关系上,它和单塔 Transformer 的差距在噪声范围内。所以如果你的业务以头部关系为主,别指望这套东西能带来明显提升。

3. 核心机制拆解:一层 URFormer 里到底发生了什么

前面讲的是"为什么这么设计",这一节讲"具体怎么算"。我会沿着数据流走一遍:输入侧的张量组织、注意力里的关系感知、专家路由的前向、预训练目标的设计,最后给一组参数量和显存的手算过程。这一节偏工程实现,如果你只想跑起来,可以直接跳到第 4 节,但我建议还是看一眼,后面排查问题时这些细节会救你。

3.1 输入侧到底喂了哪三个张量

一条训练样本不是"一个节点",而是"一个以目标节点为中心的 k-hop 子图"。我采样的是 2-hop,第一跳每类关系最多采 16 个邻居,第二跳每类关系最多采 8 个,总量上限 512 个节点。为什么按关系类型分层采样而不是全局随机?因为异构图的度分布极度不均,全局随机采样会让某些关系在子图里彻底消失,等于这个样本白采。

喂进模型的张量有三个。第一个是节点特征矩阵 X,形状是 [B, N, d_in],B 是 batch 里的样本数,N 是子图节点数上限(不足补零,配 mask),d_in 是原始特征维度,异构节点的特征维度不一致时先在输入侧做投影对齐到 256 维。第二个是节点类型 id,形状 [B, N],用来查类型的 embedding 并加到节点表征上。第三个是边的关系类型 id,形状 [B, N, N] 或者稀疏形式 [B, E],E 是子图内的边数。

第三个张量是这套方案和常规 GNN 最大的区别所在。常规做法是把边类型当成注意力偏置(类似相对位置编码),这里更进一步——每条边的关系类型 id 会被查成完整的 256 维向量,作为该条边的"关系 token"参与到后续计算里。也就是说,一条边的语义不是加在注意力分数上的一个标量偏置,而是有自己完整表征的对象。这个设计让关系类型能承载更复杂的语义,也为后面的专家路由提供了输入。

提示:边的关系类型 id 一定要用元关系 id(三元组整体),不要只用边语义 id。我第一版偷懒只用了边语义,结果(用户, 购买, 商品)和(店铺, 购买, 商品)被当成同一种关系,退货率预测的指标掉了 4 个百分点。三元组整体编码才能区分方向。

3.2 注意力计算:先关系内、再关系间

注意力层的计算分两步走,这个分层结构是保证"关系感知"不失效的关键。

第一步是关系内注意力,也就是只在同一条元关系的邻居之间做 softmax。具体地,对于目标节点 i 和它通过元关系 r 连出去的邻居 j,注意力分数是 query 与 key 的点积加上一个关系偏置项,偏置项来自 tokenizer 输出的关系向量过一个标量投影。softmax 在这条关系的邻居集合内归一化。这么做的好处是,不同关系的邻居数量差异巨大(有的关系几十个邻居,有的只有一个),如果放在一起 softmax,邻居多的关系会压倒邻居少的关系。

第二步是关系间聚合,把每条关系算出来的聚合结果当成一个新的 token,在关系维度上再跑一次注意力或者加权求和。权重由关系向量和当前节点表征共同决定,相当于让模型自己判断"对这个节点来说,哪条关系更该被关注"。这一步的计算量很小,因为一个节点的关系类型通常不超过 10 条。

实现上有个细节要注意:fp16 下做两步 softmax 容易数值溢出,尤其是在第二步关系数很少的时候。我的处理是第一步的 attention logits 减最大值(标准操作),第二步改用 fp32 计算,虽然会慢一点点,但避免了 loss 突然变 NaN。这个坑我在第 5 节还会提到。

3.3 专家路由的前向过程

节点表征经过注意力层之后进入 FFN,这里就是 MoT 生效的地方。具体流程是:拿当前 token 所属的元关系向量(如果是一个聚合后的节点 token,就取它所有入边关系的加权平均),送进路由网络得到专家 logits,取 top-2,softmax 归一化得到权重,然后分别过两个专家,输出按权重加权求和。

路由网络本身很轻,就是一层 256→专家数的线性层,参数量可以忽略。它输入的除了关系向量,我还额外拼了一个"层标识"embedding,因为不同层关注的粒度不同——浅层更关注局部结构,深层更关注语义。实测加这一项在长尾关系上有小幅提升。

训练时这个路由是可微的(softmax 加权而非硬选),推理时可以走硬选加速。有一点要强调:专家是加在 FFN 上的,不是替换 FFN。也就是输出 = FFN(x) + Σ w_k · LoRA_k(x),其中 LoRA_k 的 B 矩阵初始化为零。这样初始状态下模型行为和没有专家的 Transformer 完全一致,训练是平滑演进的。如果直接把 FFN 替换成专家,初始化那一刻分布就偏了,收敛会慢很多。

3.4 预训练目标:图本身给的监督信号

这套方案能"预训练一次,到处用",靠的是自监督任务设计。用的是最朴素的一个:给定源节点和关系类型,预测目标节点的类型。比如从"用户 u 通过'购买'关系连到了 v",模型要预测 v 属于"商品"这个节点类型。这个任务不需要任何人工标注,只要有图结构就能采样出无穷多的训练样本。

为什么选这个任务而不是链接预测?因为链接预测需要负采样,而异构图的负采样极难做对——随机采一个节点作为负例,大概率是类型不对,模型直接靠类型判断就能拿高分,学不到结构信息。节点类型预测没有这个问题,它的难度天然来自于"不同关系类型下的邻居分布差异"。

训练时会对目标节点的类型做标签平滑(label smoothing 0.1),一方面是防止过拟合,另一方面是因为有些节点的类型本身就有歧义(比如"店铺"和"品牌"在部分数据里混用)。我实测标签平滑在这个任务上带来了 1.5 个点的 Macro-F1 提升。

3.5 参数量与显存的手算过程

为了让你有直观感受,我按自己的实际配置算一遍。隐藏维度 256,4 层 Transformer,专家数 8,top-2 路由,LoRA rank 8,FFN 中间层 1024。

主干参数大致是:每层的多头注意力(4 头,256 维)约 26 万,FFN 约 52 万,两项合计 78 万,4 层 312 万。节点类型 embedding 表和关系 tokenizer 的三张表加起来约 20 万(26 种元关系对应的类型数量不多)。LoRA 专家 33 万。总计约 370 万参数,模型权重 fp16 下不到 8 MB。这个体量非常小,意味着你可以把它部署在很边缘的机器上,甚至 CPU 推理都能接受。

显存的大头在激活值而不是参数。一个 batch 32 个样本、每个样本子图上限 512 个节点、256 维 fp16,单层激活约 32×512×256×2 字节 ≈ 8.4 MB,乘上中间变量数量(约 6 倍)和层数(4 层,但要存反向传播的中间值,实际约 2 倍)后大概 400 MB。加上注意力矩阵 32×512×512×2 字节 ≈ 16.8 MB(如果不用稀疏掩码会更多),总共 1 GB 以内能跑起来。我用单卡 24 GB 训练,batch 开到 128 也没有 OOM,实际瓶颈在数据加载而不在显存。

这个账算下来得到的结论是:这套方案的工程门槛比看起来低得多。它不是一个必须大算力才能玩的东西,小团队完全跑得起。真正的成本在数据侧的 schema 整理和采样逻辑上,而不是模型本身。

4. 复现路径:从 schema 定义到零样本评测

这一节是可直接抄的部分。我按实际动手顺序写:先定义 schema 和数据集,再给采样和 tokenizer 的关键代码,然后是训练配置,最后是怎么设计零样本评测协议。每一步我都标注了容易出错的点和我的实际数值。

4.1 第一步:把 schema 定义清楚

schema 定义是整个复现过程里最容易被跳过、但影响最大的一步。我用一个 Python 字典来描述,节点类型、边语义、元关系三段分明。关键点是元关系不要手写,从节点类型和边语义自动展开——因为后面采样和 tokenizer 都要按元关系 id 索引,手写容易漏。

# schema 定义:节点类型、边语义、自动展开的元关系 NODE_TYPES = ["user", "item", "shop", "category", "brand", "coupon"] EDGE_SEMANTICS = ["click", "favorite", "cart", "buy", "refund", "belong_to"] # 每条边语义允许的 (src_type, dst_type) 组合,不写全,只写业务上真实存在的 ALLOWED = { "click": [("user", "item"), ("user", "shop")], "favorite": [("user", "item")], "cart": [("user", "item")], "buy": [("user", "item"), ("shop", "item")], "refund": [("user", "item")], "belong_to":[("item", "category"), ("item", "brand"), ("coupon", "shop")], } META_RELATIONS = [] for sem, pairs in ALLOWED.items(): for src, dst in pairs: META_RELATIONS.append((src, sem, dst)) META2ID = {mr: i for i, mr in enumerate(META_RELATIONS)} print(len(META_RELATIONS)) # 本示例展开后是 11 条元关系

这里有个容易踩的坑:ALLOWED里只写真实存在的组合,不要图省事用节点类型两两笛卡尔积。原因是你如果在 schema 里声明了 1200 种元关系,但实际只有 11 种有边,那么 tokenizer 的三张表会有大量永远不被激活的行,训练时它们拿到的是纯噪声梯度,会干扰路由网络的学习。宁可少声明,用到再加。

另外一个建议是给每个元关系加一个频次统计,训练时对极低频元关系(比如边数少于 100)做下采样或者干脆不进训练集,留给零样本评测用。这一步我在 4.4 会展开。

4.2 第二步:子图采样与 tokenizer 实现

采样逻辑我写成按元关系分层的形式。核心思路是:给定一个目标节点的元关系集合,对每种元关系独立采样固定数量的邻居,然后拼成一个子图。这样能保证子图里的关系类型分布相对均衡,不会被超级节点主导。

import random def sample_subgraph(seed_node, adj_by_meta, meta_ids, fanout=(16, 8), max_nodes=512): """adj_by_meta: {(meta_id, node_id): [neighbor_ids]},按元关系分桶存储""" seen = {seed_node} frontier = [seed_node] node_list, meta_list, src_idx, dst_idx = [], [], [], [] for hop, k in enumerate(fanout): next_frontier = [] for u in frontier: for mid in meta_ids: nbrs = adj_by_meta.get((mid, u), []) if not nbrs: continue # 采样 k 个邻居,如果邻居数远大于 k 则不重复采样 picked = nbrs if len(nbrs) <= k else random.sample(nbrs, k) for v in picked: if len(node_list) >= max_nodes: break if v not in seen: seen.add(v) next_frontier.append(v) node_list.append(v) src_idx.append(u) dst_idx.append(v) meta_list.append(mid) frontier = next_frontier[:64] # 控制二跳的入口数量,防止爆炸 if not frontier: break return node_list, meta_list, src_idx, dst_idx

这段代码里frontier = next_frontier[:64]那一行是我调了两天才加的。不加的话二跳入口可能上千个,采样耗时直接失控,而且采出来的节点 90% 都在同一个超级节点的邻域里,信息量极低。截断到 64 个入口之后,采样耗时从每 batch 1.2 秒降到 0.2 秒,指标基本没变。

tokenizer 的实现就是三段查表拼接加投影:

import torch import torch.nn as nn class RelationTokenizer(nn.Module): def __init__(self, num_src, num_sem, num_dst, dim=128, out_dim=256): super().__init__() self.emb_src = nn.Embedding(num_src, dim) self.emb_sem = nn.Embedding(num_sem, dim) self.emb_dst = nn.Embedding(num_dst, dim) self.proj = nn.Linear(dim * 3, out_dim) self.norm = nn.LayerNorm(out_dim) # 三段可学习重要性权重 self.weights = nn.Parameter(torch.ones(3)) def forward(self, src_ids, sem_ids, dst_ids): w = torch.softmax(self.weights, dim=0) e = torch.cat([ self.emb_src(src_ids) * w[0], self.emb_sem(sem_ids) * w[1], self.emb_dst(dst_ids) * w[2], ], dim=-1) return self.norm(self.proj(e))

注意self.weights那个乘法的位置。我一开始是在拼接之后做加权求和,效果反而变差了,因为拼接后的三维度之间本来就有量纲差异,加权求和会放大波动。改成先乘再加投影,训练曲线明显更稳。这个小改动带来的提升在长尾关系上有 1 个多百分点。

4.3 第三步:训练配置与超参

训练配置我整理成了一张表,这些都是我在小规模基准上试出来比较稳的值。规模更大的时候需要按比例调整,尤其是学习率和 batch size。

配置项取值说明
优化器AdamWβ2 用 0.95 而不是 0.999,图类任务上更稳
主干学习率2e-4预训练时用 1e-4 也可以,预训练更长就调小
专家(LoRA)学习率1e-3专家的量级小,学习率可以放大 5 倍
路由学习率5e-4介于两者之间,太大容易坍缩
warmup总步数的 5%路由网络特别需要 warmup
batch size64显存允许就往上加,128 也跑得动
梯度裁剪1.0必开,专家梯度偶尔会尖峰
dropout0.1子图内部,节点特征上 0.1 就够
负载均衡损失系数0.01关键参数,第 5 节细说
预训练轮数15 到 20用验证集节点类型 Macro-F1 早停

启动命令很朴素,用标准的分布式训练就行:

python -m torch.distributed.run --nproc_per_node=2 train_pretrain.py \ --schema config/schema_demo.json \ --graph_root /data/global_hin \ --hidden 256 --layers 4 --heads 4 \ --num_experts 8 --topk 2 --lora_rank 8 \ --batch 64 --lr_main 2e-4 --lr_expert 1e-3 \ --balance_coef 0.01 --epochs 20

主干和专家的分离学习率必须分开设,这是我在复现时最容易忽略的一点。如果共用同一个学习率,要么主干学得不够,要么专家过拟合到少数几个高频关系上。我用的是 parameter group 分别指定,代码里按参数名前缀匹配即可。

4.4 第四步:零样本评测协议怎么设计

这一节可能是最容易被做错的地方。很多复现的指标虚高,问题都出在评测协议上。我的做法是按元关系整体留出,而不是按样本随机留出。具体地,从 26 条元关系里挑 5 条,这 5 条的所有边在训练阶段完全不出现,测试时只在这些边对应的样本上评。

挑哪 5 条也有讲究,不能随机挑。我按两个维度挑:一是要有 2 条是"组合新但零件旧"的(三段里每一段都在训练集出现过),二是要有 1 到 2 条是"零件也新"的(至少有一段完全没出现过)。这两种情况的难度差很远,混在一起报指标会误导。

实际数值上,我实测的结果大致是这样:组合新的元关系上,零样本的节点类型预测 Macro-F1 能到完全监督版本的 82% 左右;零件也新的那种,只有 63% 左右,基本就是靠类型 embedding 的初始化和路由先验在撑。这个落差需要在方案宣传时如实说明——"任意关系类型"指的是架构上支持,不是指标上无损。

还有一个细节:零样本评测时,新元关系的 embedding 是随机初始化的。如果你能拿到新关系的少量边(比如 50 条),做几步微调(只训 tokenizer 和门控,冻住主干),指标能拉回到完全监督的 90% 以上。我把这个叫"热启动",实际落地时这个设置更现实——纯零样本更多是研究指标,业务上一般多少能拿到一点样本。

4.5 实测的吞吐与成本

最后给一组实际数字,方便你判断资源需求。在 120 万元关系边、平均子图 180 个节点的配置下,双卡 24 GB 训练,一个 epoch 大约 22 分钟,20 个 epoch 跑完约 7.5 小时。推理侧,单卡 batch 256 的 QPS 约 1400(fp16),如果开启专家硬选和 fp32 只在关系间聚合那一步用,QPS 能提到 1800 左右。

这套数字意味着,一次完整实验的算力成本是很低的。真正的时间成本在数据侧——把业务图整理成统一的 HIN 格式、给每条边打上正确的元关系 id、处理脏数据里的类型不一致,这部分我花了整个项目六成的时间。如果你准备动手,建议先把 schema 和数据清洗做扎实,模型侧反而是最快的一步。

5. 实战问题排查:几个把我卡住过的坑

理论和代码都写完了,实际跑起来才是见真章的时候。这一节我把遇到过的问题、定位思路和解决方式整理出来,其中前两个是这套方案独有的(和专家路由有关),后面几个是异构图模型的通病,但在这个框架下表现有差异。

5.1 专家路由坍缩:8 个专家最后只有 2 个在干活

这是最常见也最隐蔽的问题。表现是训练 loss 正常下降,验证指标也不算差,但你去看路由的专家使用统计,会发现 8 个专家里有 5 个的调用次数接近于零,所有 token 都挤在两三个专家上。这个时候模型退化成了一个参数量更大的普通 Transformer,专家机制形同虚设。

原因是路由网络的正反馈循环:某个专家早期因为随机初始化恰好表现稍好,拿到的梯度就更多,表现更好,然后被选得更多。这个马太效应在训练初期几轮内就会固化。

解决办法有三层。第一层是必加的负载均衡损失,形式是专家使用频率的方差项,系数 0.01。这个系数很敏感,0.001 基本没效果,0.1 会压得路由无法区分关系类型,指标反而掉。第二层是路由 logits 加噪声,训练时给 logits 加一点高斯噪声(标准差 0.1),推理时去掉,这个技巧能显著缓解早期固化。第三层是监控,我会在训练日志里每 100 步打印一次专家使用分布,一旦发现某个专家连续 500 步零调用,就说明均衡没做好,得回去调系数。

注意:不要用"硬性轮转"(强制轮流选专家)来解决坍缩。我试过,训练 loss 能降但验证指标明显变差,因为强行让不合适的专家处理不熟悉的关系类型,等于给模型注入了错误先验。均衡损失加噪声是更温和也更有效的路子。

5.2 新关系类型上线后指标反而下降

这个问题比路由坍缩更让人抓狂:明明是为了支持新关系,结果加进去之后老关系的指标也掉了。我遇到过两次,原因不同。

第一次是词表污染。新加的元关系在 tokenizer 的三段里引入了新的节点类型,而这些新类型的 embedding 是随机初始化的,训练时它们参与到了共享的投影层和路由网络里,把原本稳定的分布带偏了。解决办法是给新加入的类型 embedding 加一个小的初始化方差(比如 0.02 而不是默认的 1.0),并且在头几百步冻结投影层,只训新表。

第二次是采样预算被挤占。全局图的子图节点上限是固定的(我设的 512),新增关系之后,同样的预算要分给更多关系,导致每条关系分到的邻居数变少,老关系的子图信息变稀疏了。这个问题的定位方法很直接:对比加关系前后每个样本的平均子图节点数和每种关系的平均邻居数,一看就明白。解决办法是把上限提到 768,或者在采样时按关系类型设优先级,高频关系保持原配额。

5.3 一份可以贴在显示器上的排查表

我把常见现象和对应原因整理成表,方便你按症状查。

现象最可能的原因排查动作处理方式
loss 突然变 NaN关系间聚合用了 fp16打印该步的 logits 范围关系间 softmax 强制 fp32
验证指标卡在基线不动专家全坍缩或全不用看专家调用分布调均衡系数到 0.01,加路由噪声
长尾关系指标极差采样时被高频关系挤掉统计每种关系的平均邻居数按关系分层采样,设最小配额
零样本新关系几乎随机三段里有新零件检查新元关系三段是否都见过加少量样本热启动微调
训练快但推理慢路由走了 soft 加权profile 路由耗时占比推理改硬选 top-2
加新关系后老关系掉点子图预算被摊薄对比加前后的子图规模提上限或设关系优先级
指标波动大,重跑差异明显采样随机性太强固定种子重跑三次固定采样种子 + 增大 batch

这张表里我觉得最值得记的是"固定种子重跑三次"这一条。图模型的方差本来就大,很多所谓的"改进"其实是随机波动。我现在的习惯是任何改动都要跑三个种子取均值,差异小于 0.5 个百分点的一律当噪声处理,不写进结论。

5.4 几条不太会写在文档里的心得

第一条:先把单关系跑通再上全局图。我一开始就直奔全局异构图的实现,结果采样、类型对齐、路由三处问题同时出现,根本没法定位。后来退回单场景(只有 6 种元关系)把整条链路跑顺,再逐步加场景,问题就一个一个浮出来了。这个顺序看起来慢,实际快得多。

第二条:关系类型词表要留扩展位。embedding 表可以预留 20% 的空行,方便后续加新类型时不用改 checkpoint 结构。我第一版表大小写死了,后来加关系只能重训,白扔了两天。

第三条:别迷信 top-k 路由的 k 值。很多人觉得专家多了 k 就该大,其实不然。我的实验里专家数从 4 加到 8 时,top-2 一直是最优;加到 16 个专家时 top-2 仍然够用,因为路由本来就是在做稀疏选择,k 大了反而破坏了稀疏性带来的正则效果。

第四条:评估时一定要分关系类型看指标,不能只看总体均值。总体均值会被高频关系主导,把长尾关系的恶化完全掩盖。我习惯在评估脚本里输出一张按元关系分组的指标表,每次看这张表比看总均值有用十倍。

我个人在实际操作中的体会是,这类"统一建模"方案的真正价值不在于论文里那个漂亮的平均指标,而在于它把"支持新关系类型"从一次工程改造变成了一次配置改动。这个转变对长期维护的意义,远大于单点指标的提升。至于后续还能怎么扩,我目前在试两个方向:一是把关系类型的层级结构(比如"购买"是"互动"的子类)编码进 tokenizer,二是让专家在层与层之间共享而不只是层内独立,前者看起来更有戏。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进