最近在做一个药物重定位相关的实践项目,最初的思路很朴素:用传统的分子结构相似性去预测药物和靶点之间的关联。但做了几轮后,问题很明显——结构相似性只能捕捉化学层面的特征,完全忽略了隐藏在图谱背后的生物学语义关系。于是,我换了方向,用知识图谱把药物、靶点、疾病、副作用之间的关系重新建模,再配合图神经网络做表示学习,最后用消融实验验证每个设计模块的贡献。这篇文章会把这套方案的完整实现流程、模型设计、消融实验过程和实际踩坑经验都拆开来讲,代码部分基于PyTorch和常见图计算库,适合有一定Python基础、想系统了解知识图谱在药物靶点预测中如何落地的开发者参考。
1. 项目整体设计与思路拆解
1.1 为什么是知识图谱而不是传统分子特征
药物靶点相互作用预测不是新问题,过去的主流做法是分子对接、序列比对、QT建模,或者用分子指纹加传统分类器。这些方法的通病是:特征空间过于局部化。分子指纹只看原子和化学键的局部模式,序列比对只看氨基酸片段的保守性,它们都不会告诉你“这个药物恰好作用于某条通路的中间节点”这样高层次的关系。
知识图谱的建模方式完全不同。我拿到一组三元组,比如Drug_A - binds_to - Target_B、Target_B - participates_in - Pathway_C、Drug_A - treats - Disease_D,这些关系组合起来,就把原本孤立的生物实体放进了上下文丰富的网络里。做预测时,不再只依赖药物本身是什么,而是依赖它在一张图里的位置、邻域结构、与哪些靶点共享了疾病通路。这种语义增强的表示,对新靶点、老药新用的场景尤其有帮助。
1.2 消融实验在设计阶段就必须想清楚
很多人在做知识图谱类的预测模型时,喜欢把能加的模块全堆上去——全局邻居采样、多关系编码、属性嵌入、辅助损失,最后模型效果确实好,但问哪个模块贡献了什么,完全说不清。所以我在项目设计的第一天就把消融实验写进了开发计划:模型必须模块化,每个模块能单独开关,每一处设计决策都有对应的“对照组”。
这样做的实际价值不止是为了论文里那张表。更重要的是,消融实验能帮你抓到模型“偷懒”的行为。比如,有时候去掉全局邻居信息之后效果反而更好,这说明全局聚合层在引入噪声而不是信息;又比如,辅助疾病嵌入的剔除导致AUC骤降,说明模型依赖的恰恰是疾病通路信息,而不是药物本身。这些发现都会直接影响下一轮迭代方向。
1.3 技术选型与架构模式
整个系统的技术栈选了PyTorch作为深度学习框架,图的存储和批采样用常见图数据工具,模型层自己实现图消息传递。选择PyTorch而不是TensorFlow,主要考虑是调试灵活度更高,自定义GNN层和loss更顺手;图工具用来做三元组加载、邻域采样、批次划分,可以让数据管道的开发量大幅减少。
整体流程如下:
- 从公开生物数据库抽取药物、靶点、疾病、通路等实体和关系,清洗成三元组格式。
- 对三元组做约束校验,过滤缺失实体和孤立节点。
- 对实体和关系做ID映射,构建统一的整数索引。
- 设计负采样策略,生成不存在的三元组作为反例。
- 构建GNN模型,对每个实体生成嵌入向量。
- 用评分函数计算药物-靶点对的得分,二分类任务训练。
- 保存最优模型,在测试集上评估,并执行多组消融实验。
2. 数据准备与知识图谱构建
2.1 实体与关系的三元组设计
知识图谱的质量直接决定模型上限。我最初从公开数据库抽取数据时,图里混了多种实体类型,包括药物、靶点、疾病、通路、副作用、适应症等。实体类型太多会让模型参数爆炸,所以我最终保留了四类核心实体:
- 药物节点,约4000个;
- 靶点节点,约2500个;
- 疾病节点,约3300个;
- 通路节点,约980个。
关系方面,设计了8种有明确生物学含义的类型,包括binds_to、treats、involved_in_pathway、causes_side_effect、associated_with_disease、expressed_in_tissue、belongs_to_class、interacts_with。
三元组的设计原则有三条。首先,所有关系都要求有来源或文献依据,不能凭空构造;其次,每个实体至少关联两个三元组,避免孤立节点;最后,关系方向要严格定义,比如binds_to的方向必须是药物指向靶点,不能随意颠倒。
2.2 负样本生成策略
知识图谱版本图数据里的正样本是有明确依据的真实三元组,负样本则完全不存在于图中。最难的一步是怎么生成“看起来像正样本但实际为假”的替代性三元组。如果负样本太简单,比如随机把药物和完全不相关的实体组合,模型会通过记忆实体频率来偷懒,精度虚高但没有泛化能力。
我采用的是混合负采样策略:80%的负样本用随机替换头实体或尾实体的方式生成,20%用难度更高的近邻替换方式,替换成与当前实体嵌入距离接近的候选实体。这样既保留了随机采样的多样性,又强迫模型学会区分细微差异。
一个关键细节是,负样本必须做去重检查。替换头实体或尾实体后,新三元组有可能误打误撞出现在训练集的正样本里,如果不加过滤,会造成标注冲突,模型信号直接混乱。我在构建数据管道的代码里专门加了一个集合判断步骤:
import random def generate_negative_samples(positive_triples, all_entities, num_negatives=5): """ 混合负采样:随机替换 + 邻域替换 需要保证生成的三元组不在正样本集合中 """ positive_set = set(('|'.join(map(str, t)) for t in positive_triples)) negatives = [] # 先构建实体类型的合法映射:每个关系位点允许哪些实体类型 entity_type_dict = {} for h, r, t in positive_triples: entity_type_dict.setdefault((h, r), set()).add(t) entity_types = {} for h, r, t in positive_triples: entity_types[h] = h[0] # 假设h是(type_id, entity_id) entity_types[t] = t[0] while len(negatives) < num_negatives: h, r, t = random.choice(positive_triples) mode = random.random() if mode < 0.8: # 随机替换 if random.random() < 0.5: candidate = random.choice(all_entities) new_triple = (candidate, r, t) else: candidate = random.choice(all_entities) new_triple = (h, r, candidate) else: # 近邻替换(简化版:在相同类型的实体池中随机选一个 # 真实场景可以基于embedding距离选择近邻实体) candidate = random.choice( [e for e in all_entities if e[0] == t[0]] ) new_triple = (h, r, candidate) key = '|'.join(map(str, new_triple)) if key in positive_set: continue negatives.append(new_triple) return negatives负采样比例我用的是1:5,每个正样本对应5个负样本。这个比例不是拍脑袋定的。试过1:1,模型精度不足,因为正负样本极度不平衡,模型倾向将所有样本预测为负类;试过1:10,训练时间拉长了一倍,但AUC提升幅度很小,边际收益递减。1:5是在测试集上表现稳定且训练开销相对平衡的选择。
2.3 批次采样与图邻域构建
训练GNN时不可能每次全图加载。我按药物-靶点对做批次采样,每个批次内部包含若干组药物-靶点关系对及其对应的邻域子图。邻域采样的思路和GraphSAGE类似,选定中心节点后,在固定层数内做邻居抽样,每层采样固定数量的邻居节点。
需要注意的是,药物和靶点在图里的邻居构成完全不同。药物的邻居多为靶点、疾病和副作用,靶点的邻居则多为药物、通路和相互作用靶点。为了让初始embedding与实体类型关联,我对不同类型的实体用了不同的初始特征投影层。这样药物和靶点即使在拓扑结构上很相似,其语义含义也能保留差异。
3. 模型核心实现与训练细节
3.1 图卷积信息聚合层
模型的核心是信息传递层。我采用的聚合方式如下:
h_u^{k+1} = ReLU(W_k * AGG({h_v^k : v 属于 N(u)}))在聚合函数的选择上,我对比了Mean聚合、Max聚合和注意力聚合。Mean聚合实现最简单且训练稳定,适合作为基线;Max聚合对高响应邻居更敏感,但在稀疏图中容易受单点噪声影响;注意力聚合理论上最灵活,实际调参坑最多,温度参数稍不合适权重就集中到少数几个节点上。
最终选型是Mean聚合加一层注意力偏置。具体实现是对不同邻域节点加上类型可变的加权系数,比如通路节点的权重比副作用节点高,因为通路共现对靶点预测的生物学语义更强。
模型完整结构分四层:
- 实体类型投影层,把不同类型节点的初始特征压到统一维度;
- 两层图卷积,节点特征经过两跳邻居的信息聚合;
- 关系感知融合层,根据关系类型对节点对做组合特征;
- 评分输出层,把组合特征映射成预测分数。
3.2 评分函数的选取
药物-靶点预测的本质是判断给定三元组(h, r, t)是否成立。我试了三种评分方式,差异非常明显。
第一种是点积评分,直接计算药物嵌入与靶点嵌入的内积。优点是计算快、解释性强,缺点是表达能力有限,只捕捉线性相关性。
第二种是DistMult风格的双线性评分,对关系矩阵做对角分解,能够捕捉“药物嵌入乘以靶点嵌入加上关系旋转”这种交互模式,在知识图谱补全任务上表现优于点积。
第三种是拼接后过MLP的评分,也就是把药物嵌入、靶点嵌入、关系嵌入拼接后喂给两层全连接网络。这种表达能力强,但参数量显著增大,训练需要更多数据。
最终我选择了双线性评分作为主评分函数,同时保留了MLP评分作为消融实验的对照组之一。实际结果里,双线性评分的泛化性比MLP更好,尤其在测试集上,AUC高出0.04左右。原因也好理解,药物-靶点关系在图上本身就带有一定的低秩结构,双线性形式恰好匹配了这种结构。
3.3 训练策略与损失函数
训练阶段我用的是标准的二分类二元交叉熵损失,对每个批次计算正样本三元组和负样本三元组的得分,然后套用sigmoid加BCE。公式上很简单,但有几个细节直接影响最终效果:
- 用
label_smoothing把正样本的标签平滑到0.95左右,抑制过拟合; - 对负样本分数加了较大的L2正则权重,避免所有实体嵌入范数退化为0;
- 使用Adam优化器,初始学习率0.001,训练过程中每5轮乘以0.8衰减;
- 设置早停机制,验证集AUC连续10轮不涨就停止训练,并在验证集最优时保存模型参数。
训练过程比较典型,前20轮loss快速下降,中间阶段缓慢收敛,到第60轮左右开始出现过拟合迹象,早停触发在第72轮。验证集上最终的AUC稳定在0.91左右,PRC稳定在0.87左右。
关键训练代码示例:
import torch import torch.nn.functional as F def train_step(model, optimizer, batch, device): """ batch: 包含药物索引、靶点索引、标签、邻域信息 """ model.train() optimizer.zero_grad() drug_emb = model.forward_drug(batch['drug_index'], batch['drug_neighbors']) target_emb = model.forward_target(batch['target_index'], batch['target_neighbors']) score = model.score(drug_emb, target_emb, batch['rel_type']) pos_mask = batch['label'] == 1 neg_mask = batch['label'] == 0 # 平滑正样本标签 smoothed_pos = torch.full_like(score[pos_mask], 0.95) loss = F.binary_cross_entropy_with_logits(score[pos_mask], smoothed_pos) loss += F.binary_cross_entropy_with_logits( score[neg_mask], torch.zeros_like(score[neg_mask]) ) # 嵌入正则 emb_reg = model.embedding_l2_norm() loss += 1e-4 * emb_reg loss.backward() optimizer.step() return loss.item()3.4 特征标准化与批次稳定
图数据有个容易忽略的问题:不同节点的邻域数量差异非常大。有的药物节点在途中有几百个邻居,有的只有几个。直接做邻居平均会导致度数高的节点embedding被过度平滑,度数低的节点信息匮乏。我的做法是做了邻居数量的截断和归一化,每个节点最多采样40个邻居,不足的用零向量补齐,同时记录掩码,聚合时对mask处理,避免填充值污染结果。
还有一个值得提的细节:实体ID的映射顺序会影响随机初始化embedding的分布,进而影响最终的复现性。我在预处理阶段固定了实体索引的随机种子,并且把所有数据集的划分方式和ID映射统一暴露在配置文件中,保证每次实验之间只有单一变量在变化。
4. 消融实验的设计与执行
4.1 消融维度拆解
消融实验是整个项目的核心验证环节。我设置了四个维度,每个维度分别关掉一个模块,并设置一个“完全版模型”和一个“极简版基线模型”作为上下界参考:
- 完整模型:包含实体类型投影、全局疾病特征、PO通路信息、双线性评分、混合负采样;
- 无全局疾病特征:去掉疾病节点的全局统计特征,只保留局部图结构;
- 无通路信息:删除
involved_in_pathway关系,观察通路语义的贡献; - 无关系感知:评分时忽略关系类型,模型无法区分
binds_to和treats; - 无邻域采样:每个节点只使用自身特征和单跳邻居,不扩展多跳信息;
- 极简基线:只用药物和靶点的ID映射嵌入做点积预测,无任何图结构信息。
每组实验在相同划分的数据集上训练,采用相同随机种子,保证唯一变量是被消融的模块。这个设计原则非常重要,否则消融结果无法归因。
4.2 消融实验的执行逻辑
代码层面,我采用了配置驱动的执行方式。每轮消融实验传入不同的config文件,模型根据配置决定是否启用对应模块,训练脚本完全复用。这样一个实验矩阵跑起来很轻松,不需要为每个变体写单独的训练代码。
消融执行脚本的骨架如下:
configs = [ {"name": "full", "use_global_disease": True, "use_pathway": True, "use_rel_type": True, "neighbor_depth": 2}, {"name": "no_global_disease", "use_global_disease": False, "use_pathway": True, "use_rel_type": True, "neighbor_depth": 2}, {"name": "no_pathway", "use_global_disease": True, "use_pathway": False, "use_rel_type": True, "neighbor_depth": 2}, {"name": "no_relation", "use_global_disease": True, "use_pathway": True, "use_rel_type": False, "neighbor_depth": 2}, {"name": "no_multi_hop", "use_global_disease": True, "use_pathway": True, "use_rel_type": True, "neighbor_depth": 1}, {"name": "baseline", "use_global_disease": False, "use_pathway": False, "use_rel_type": False, "neighbor_depth": 1}, ] for config in configs: model = build_model(config) run_training(config, model) eval_and_save(config, model)这个配置驱动方式的另一个好处是,实验记录的格式天然统一。每次运行结束后,会把模型配置、参数数量、训练轮数、最优验证指标、测试指标全部写入一个汇总文件,后续整理实验结果时不需要手动拼接多个日志。
4.3 评估指标的选择
药物靶点预测是典型的类别不平衡任务,正样本占比通常不足20%,单纯看准确率没有意义。我选用了三组核心指标:
- AUC,衡量排序能力,对样本比例不敏感;
- PRC,精确率召回率曲线下面积,适合正样本稀少的场景;
- 召回率@Top50,重点关注模型能否在预测分数最高的50个候选里覆盖更多真实靶点,这个指标对药物重定位的实际应用更具参考价值。
评测时,我将数据集按7:2:1划分为训练集、验证集和测试集,划分操作在预处理阶段完成后就固定下来,测试集不参与任何调参。
我同时计算了在不同负采样比例下训练出的模型在测试集上的指标变化。负采样比例从1:1调整到1:10的过程中,AUC先升后降的曲线非常典型,最高点在1:5附近,这也验证了之前关于负采样比例的判断。实验记录表格如下:
| 实验配置 | AUC | PRC | 召回率@Top50 |
|---|---|---|---|
| 完整模型 | 0.917 | 0.874 | 0.63 |
| 去掉全局疾病特征 | 0.881 | 0.832 | 0.54 |
| 去掉通路信息 | 0.862 | 0.801 | 0.47 |
| 去掉关系感知 | 0.893 | 0.845 | 0.58 |
| 去掉多跳邻域 | 0.905 | 0.853 | 0.60 |
| 极简基线 | 0.742 | 0.655 | 0.31 |
4.4 消融结果的解读逻辑
逐项分析后,一个有意思的结论浮现出来:通路信息的贡献比重比原先预期的更大。去掉involved_in_pathway关系之后,AUC下降了5.5个百分点,是所有模块中影响幅度最大的。这说明在图数据中,通路语义充当了连接药物和靶点的关键桥梁。很多药物和靶点之间没有直接的联系边,但它们在通路层面共享参与关系,GNN通过多跳消息传递把这种间接关联编码进了节点表示。少了这一层,等同割断了图的骨干结构。
全局疾病特征的影响也值得重视。在这个数据集中,同一个疾病相关的药物-靶点对往往有较高的共现度,疾病特征为节点嵌入提供了额外的聚类信号,去掉后AUC掉了3.6个百分点。
关系感知的贡献相对较小,只有2.4个百分点的降幅。我猜测原因是:当前的实体类型分布下,关系和实体类型高度共线,比如binds_to关系的尾实体几乎全是靶点实体,去掉关系类型后,模型依然可以从实体本身的类型识别中学到类似信息。
多跳邻域的贡献也比较微弱。原因是数据集里药物-靶点正例本身大多存在直接或一阶邻接关系,多跳信息更多起到平滑作用,而非根本性的信息补充。这个发现有助于针对性地优化模型结构——如果后续要提升多跳信息的有效性,可能需要采用更复杂的路径编码方式,而不是单纯加深邻居层数。
5. 常见问题与排查技巧实录
5.1 模型不收敛且Loss震荡
第一个遇到的典型问题是,在使用混合负采样后,训练loss一直呈锯齿状震荡,收敛速度明显变慢。排查后定位到两个原因。
原因一是负样本的难度差异太大。随机替换生成的负样本很容易被区分,loss快速下降后接近0;而邻域替换生成的负样本难度非常高,模型持续犯错,导致loss反复回升。处理方法是分别统计两类负样本的损失,并按权重动态调整难度比例。我后来把训练早期的高难度负样本比例降到10%,随训练轮数逐步提高到20%。这样模型先学会简单区分,再逐步适应细微差别。
原因二是批次内的邻居采样在每次epoch中都完全随机变化。同一个药物节点在不同epoch的批次里采样的邻居集合不同,导致每次前向计算时节点嵌入的上下文漂移,模型优化目标不稳定。固定采样种子后,loss曲线明显平稳。
5.2 验证集AUC高但测试集结果差
这类现象的根源在数据划分。如果药物在训练集和测试集中存在重叠,即同一个药物在不同实体对中分别出现,测试指标会虚高。我在初始划分时没有强制按药物去重,导致测试集看起来效果好,但一旦按药物维度重新划分,AUC立刻掉了近7个百分点。
修正方式是做冷启动评估:保证测试集中出现的药物在训练集中完全没有出现过。这样的划分方式虽然让指标更“难看”,但更贴近真实应用的场景——新药的靶点预测问题。这也是知识图谱预测任务里最常见的评估陷阱之一。
5.3 图特征和多模态特征融合不匹配
模型后面加入过一个外部特征融合分支,把药物的分子指纹特征编码后,与GNN的拓扑嵌入拼接。但实验发现融合后AUC反而低于纯GNN。原因是分子指纹的维度为1024维,而GNN嵌入维度只有128维,拼接后分子指纹特征在数值范围上也高出许多,GNN嵌入的表达被稀释。
解决办法不是盲目调维度,而是先做了两种特征的标准化,然后用一个门控融合层学习两个特征空间的动态权重。门控融合的加入大概稳定提升了2到3个百分点的AUC,而且让模型对特征维度的敏感度下降了不少。
5.4 常见问题速查表
| 现象 | 原因 | 解决方案 |
|---|---|---|
| Loss震荡剧烈 | 负样本难度分布不均 | 降低高难度负样本初始比例,逐步提升 |
| 验证集AUC高但测试结果差 | 训练测试划分存在实体重叠 | 按药物维度做冷启动划分 |
| 图嵌入与外部特征融合无效 | 特征维度和数值范围差异大 | 标准化后加门控融合层 |
| 邻居采样导致复现困难 | 批次内采样随机变化 | 固定采样种子或缓存邻居集合 |
| 过度平滑,节点嵌入趋同 | 邻居数量不均衡,层数过深 | 限制邻居采样数,归一化聚合 |
| 模型预测分数整体偏低 | 正样本标签被严格设为1 | 使用标签平滑,设置为0.9-0.95 |
5.5 关于随机种子与复现的深度实践
生物信息领域的图神经网络实验,复现性往往比视觉方向的实验更难,原因是图数据的计算路径高度依赖邻居采样顺序和批次构造顺序。我踩过好几次坑,最终形成了一套标准流程:
- 在数据预处理阶段固定所有随机种子;
- 对邻居采样做缓存,同一个epoch内每次访问同一节点都返回相同邻居集合;
- 不同epoch之间允许邻居变化,但epoch内部的随机种子保持不变;
- 所有实验共享相同的初始权重状态,即先保存一份统一的模型初始化参数,再复制到各个消融变体上。
对比实验如果没有统一初始权重,少数模块插拔的性能差异可能只是初始化的随机波动,而不是设计差异造成的。只有统一了模型初始化,消融实验的对比结果才真正可信。
5.6 训练过程中的额外观察
训练结束后,我做了一个错误样本分析,把测试集中被错误分类的样本单独导出。观察发现,错误样本主要集中在缺乏足够邻域信息的孤立靶点节点上。这类靶点在图里通常只关联一两个通路,缺乏足够的上下文信号,GNN难以学习到有效的向量表示。
针对这个现象,我尝试了用多任务学习辅助预测——在预测药物-靶点相互作用的同时,增加一个通路归属预测的辅助任务。辅助任务的梯度信号能反过来更新靶点节点表示,让孤立靶点也可以从通路关系中获得领域语义。这个改进让整体AUC又提升了1.8个百分点,尤其是在长尾靶点上的提升更明显。
项目做完后,我最大的体会是:知识图谱模型的效果上限,很大程度上不是由模型能力决定的,而是由图谱的质量和负样本的构造方式决定的。消融实验的价值也远超“写论文用”,它能强制你审视每一个设计决策的合理性,在很多场景里,结论可能跟你最初的直觉完全相反。如果你也在做类似的任务,建议把消融实验提前纳入设计流程,每个模块从第一步就保持可开关状态,后期会省下大量重复开发的时间。整套代码在本地环境运行一遍完整实验大约需要4个小时(GPU训练),建议先在小规模数据上验证模型和数据管道跑通,再上全量数据,能省去很多 debugging 的时间。