ScoreSample-SEAL：基于重要性采样的 SEAL 链路预测模型

English version is preserved below.

ScoreSample-SEAL 是一个面向链路预测与候选关系预测任务的图神经网络项目。项目基于 SEAL 框架，将节点对匹配问题建模为二分类形式的链路预测任务，并通过“重要性评分采样”改进原始 SEAL 的子图提取过程，从而减少无关邻居节点带来的噪声和计算开销。

该方法适用于推荐召回、关系推荐、候选边预测等场景，尤其关注稀疏图中候选节点对的局部结构建模。

项目特点

• 基于 SEAL 框架，将候选边预测转换为封闭子图分类任务。
• 使用 DRNL（Double-Radius Node Labeling）对节点进行结构标注，编码节点到两个目标节点的距离关系。
• 引入 ScoreGNN 作为节点重要性评分模块，用于评估候选邻居对目标节点对 (x, y) 的贡献。
• 采用自适应子图采样策略：保留高分节点，引入少量随机节点，并对过小子图进行动态扩展。
• 使用 PyTorch Geometric 完成数据切分、负采样、子图抽取、结构标注、模型训练、评估与可视化全流程。
• 在 Cora、CiteSeer、PubMed、GitHub、USAir、C.elegans 六个公开图数据集上进行实验验证。

方法概述

原始 SEAL 会为每条候选边 (x, y) 提取 n-hop 封闭子图，然后使用 DGCNN 对子图进行分类。该方法效果较好，但固定半径的子图提取容易引入与目标边关系不大的节点，尤其在局部邻域快速扩张的图中会增加噪声和计算成本。

ScoreSample-SEAL 对子图构建阶段进行了改进：

1. 对目标节点 x 和 y 提取初始 n-hop 封闭子图。

2. 如果子图节点数较少，则将半径扩展为 num_hops + gamma。

3. 如果子图较大，则使用预训练的 ScoreGNN 对候选邻居节点 v 进行重要性评分：

   s_v = (s_{v,x} + s_{v,y}) / 2
   

   其中 s_{v,x} 和 s_{v,y} 分别表示候选节点 v 与两个目标节点之间的预测关系分数。

4. 保留重要性分数排名前 alpha% 的节点。

5. 从剩余候选节点中随机采样 beta%，用于保留一定结构多样性。

6. 重建压缩后的子图，移除目标边，执行 DRNL 标注，并输入 DGCNN 进行最终链路预测。

flowchart LR
    A["输入图"] --> B["训练 / 验证 / 测试边划分"]
    B --> C["训练 ScoreGNN"]
    C --> D["候选邻居重要性评分"]
    B --> E["n-hop 封闭子图"]
    D --> F["Top alpha% + Random beta% 采样"]
    E --> F
    F --> G["DRNL 结构标注"]
    G --> H["DGCNN + SortPooling"]
    H --> I["链路存在概率"]

Loading

图示

DRNL 节点结构标注：

SEAL 与 ScoreSample-SEAL 框架对比：

模型组成

ScoreGNN：重要性评分模型

model/score_gnn.py 实现了多种 GNN 编码器：

• ScoreGCN
• ScoreGAT
• ScoreSAGE
• ScoreGIN

默认预测头为 HadamardMLPPredictor。它对两个节点 embedding 做 Hadamard 乘积，然后通过 MLP 输出节点对关系分数。

默认配置：

参数	取值
隐藏维度	256
输出维度	128
GNN 层数	3
Dropout	0.5
学习率	0.01
训练轮数	200

ScoreSample-SEAL：最终链路预测模型

最终分类器沿用 SEAL/DGCNN 思路：

• DRNL 标注目标节点与子图中的其他节点。
• GCN 层学习子图节点表示。
• SortPooling 将不同规模的子图转为固定长度表示。
• 1D 卷积与 MLP 输出最终二分类 logit。

默认配置：

参数	取值
隐藏维度	32
GNN 层数	3
SortPooling 比例 k	0.6
学习率	0.0001
训练轮数	50

重要性采样器

默认采样参数：

参数	含义	默认值
k_min	子图规模阈值	60
num_hops	初始封闭子图半径	2 或 3，依数据集而定
alpha	重要性最高节点保留比例	40
beta	剩余节点随机采样比例	20
gamma	小子图额外扩展跳数	2
score_fn	节点重要性评分函数	gnn

代码中也实现了 Adamic-Adar 和 Personalized PageRank 两种评分方式，用于消融实验。

项目结构

路径	说明
config.py	数据集、模型、采样器和训练参数配置
out_manager.py	输出目录、配置保存和日志管理工具
data_preprocess.ipynb	数据加载、训练/验证/测试划分与负采样
score_gnn.ipynb	训练 ScoreGNN 重要性评分模型
score_sample.ipynb	基于重要性分数的子图采样与 DRNL 标注
ss_seal.ipynb	ScoreSample-SEAL 主训练与评估流程
seal.ipynb	原始 SEAL 基线实现
visualization.ipynb	AUC/AP 曲线和实验结果可视化
model/score_gnn.py	ScoreGNN 编码器和链路预测头
model/ssseal.py	SEAL 风格的 GNN 分类模块
model/scoregnn.pth	已保存的 ScoreGNN 权重
data/	原始、预处理和划分后的图数据
out/	训练日志、配置文件和中间输出
result/	各数据集最终实验日志

环境依赖

项目以 Jupyter Notebook 为主要运行形式，核心依赖如下：

python >= 3.10
torch
torch-geometric
numpy
scipy
scikit-learn
networkx
tqdm
psutil
jupyter
matplotlib

安装示例：

pip install torch torch-geometric numpy scipy scikit-learn networkx tqdm psutil jupyter matplotlib

如果使用 GPU，请根据本地 CUDA 版本安装对应版本的 torch 和 torch-geometric。

运行流程

当前项目以 notebook 组织，建议按以下顺序运行。

1. 配置数据集

修改 config.py：

class Config:
    version: str = "Cora"
    dataset: str = "Cora"

支持的数据集：

Cora, CiteSeer, PubMed, GitHub, USAir, Celegans

2. 数据预处理

运行：

data_preprocess.ipynb

生成如下数据划分文件：

data/<dataset>/split/train_data.pt
data/<dataset>/split/val_data.pt
data/<dataset>/split/test_data.pt

实验中的数据划分比例为：

数据划分	比例
训练集	85%
验证集	5%
测试集	10%

正负样本比例为 1:1。

3. 训练 ScoreGNN

运行：

score_gnn.ipynb

该 notebook 会训练重要性评分模型，并保存权重：

model/scoregnn.pth

4. 生成 ScoreSample-SEAL 子图

运行：

score_sample.ipynb

生成采样后的子图数据：

data/<dataset>/split/ssseal_train_data_k<k_min>_h<num_hops>_<version>.pt
data/<dataset>/split/ssseal_val_data_k<k_min>_h<num_hops>_<version>.pt
data/<dataset>/split/ssseal_test_data_k<k_min>_h<num_hops>_<version>.pt

对于较大数据集，notebook 会先按 batch 保存正负样本子图，再进行合并，以降低单次内存压力。

5. 训练并评估 ScoreSample-SEAL

运行：

ss_seal.ipynb

训练过程中会输出：

Train_AUC, Train_AP, Val_AUC, Val_AP, Test_AUC, Test_AP

6. 运行基线与可视化

原始 SEAL 基线：

seal.ipynb

实验结果可视化：

visualization.ipynb

数据集

数据集	节点数	边数	平均度	特征维度
USAir	332	2,126	12.81	N/A
C.elegans	297	2,148	14.46	N/A
Cora	2,708	5,429	4.00	1,433
CiteSeer	3,327	4,732	2.84	3,703
PubMed	19,717	44,338	4.50	500
GitHub	4,965	18,547	7.47	N/A

实验结果

评价指标：

• AUC：ROC 曲线下面积。
• AP：平均精度。
• Time：平均每轮训练耗时。

AUC

数据集	GCN	SEAL	ScoreSample-SEAL
Cora	86.68 +/- 1.57	91.30 +/- 0.85	93.98 +/- 0.35
CiteSeer	86.01 +/- 1.72	85.61 +/- 1.55	91.10 +/- 0.65
PubMed	92.34 +/- 0.58	96.93 +/- 0.39	96.44 +/- 0.38
GitHub	88.63 +/- 1.21	91.27 +/- 0.54	91.57 +/- 0.35
USAir	95.28 +/- 0.76	96.99 +/- 0.31	97.46 +/- 0.12
C.elegans	90.72 +/- 1.32	89.50 +/- 0.63	89.87 +/- 0.44

AP

数据集	GCN	SEAL	ScoreSample-SEAL
Cora	85.62 +/- 1.18	92.83 +/- 0.78	94.53 +/- 0.22
CiteSeer	86.61 +/- 0.93	88.62 +/- 1.02	93.07 +/- 0.88
PubMed	92.16 +/- 1.21	96.93 +/- 0.53	96.54 +/- 0.61
GitHub	90.20 +/- 0.72	92.80 +/- 0.85	92.39 +/- 0.49
USAir	95.78 +/- 0.86	97.45 +/- 0.32	97.64 +/- 0.38
C.elegans	88.77 +/- 1.37	89.28 +/- 0.69	89.47 +/- 1.06

平均每轮训练耗时

数据集	SEAL	ScoreSample-SEAL
Cora	12.00 +/- 1.03	11.92 +/- 2.10
CiteSeer	9.80 +/- 1.11	8.79 +/- 1.54
PubMed	64.17 +/- 6.82	251.29 +/- 14.34
GitHub	87.48 +/- 7.74	311.23 +/- 13.13
USAir	2.92 +/- 0.48	3.11 +/- 0.42
C.elegans	2.96 +/- 0.39	2.87 +/- 0.52

实验结果表明，ScoreSample-SEAL 在 Cora、CiteSeer、GitHub、USAir 等数据集上取得了更高的 AUC 或 AP，尤其在稀疏图场景中提升明显。但在 PubMed、GitHub 等较大图上，重要性评分与动态采样会带来额外计算开销。

Cora 消融实验

模型	AUC	AP
Variant 1：用 GNN 重要性分数替代 DRNL	80.44 +/- 1.92	81.98 +/- 1.74
Variant 2：保留 DRNL，固定 k 和 n	92.09 +/- 0.61	93.23 +/- 0.42
Variant 3：固定 k，动态扩展 hop	90.50 +/- 0.78	91.91 +/- 0.69
Variant 4：Adamic-Adar 重要性评分	89.36 +/- 1.13	91.01 +/- 1.09
Variant 5：Personalized PageRank 重要性评分	93.20 +/- 0.49	94.04 +/- 0.37
ScoreSample-SEAL	93.98 +/- 0.35	94.53 +/- 0.22

消融实验说明，DRNL 结构标注对 SEAL 类方法仍然非常关键；在保留 DRNL 的基础上，使用 GNN 学习得到的重要性评分优于手工启发式评分。

局限性

• 对大规模图，动态评分和采样会显著增加计算成本。
• alpha、beta、gamma、k_min、num_hops 等超参数需要手动调节。
• 当前动态采样策略仍然是规则驱动的，灵活性有限。
• DRNL 是固定结构编码，对复杂结构模式的适应能力有限。

后续优化方向

• 优化大图上的批量子图评分与采样效率。
• 引入自动化超参数搜索。
• 探索更灵活的节点重要性评估机制。
• 改进 DRNL 之外的结构标注方式。
• 将当前 notebook 工作流重构为可复用的命令行脚本。