PENCIL 时序链接预测 — 计图挑战赛

战队：手工智能 | 基于 PENCIL (Plain Transformers) 的时序图链接预测方案

灵感来源：Plain Transformers are Surprisingly Powerful Link Predictors (Truong et al., ICML 2026)

---

项目概述

本项目实现了一个基于 PENCIL 思想的时序图链接预测模型，用于计图挑战赛赛道一。

核心思路：用一个普通的 Transformer（无需 GNN、无需节点 ID 嵌入表），通过对采样局部子图做 token 化，把图结构显式编码进 token，在时序链接预测任务上超越各种复杂的专用模型。

当前主线 (main_pencil_fast.py) 融合了：

• PENCIL Transformer：子图 token 化 + 冻结正交投影 + 邻接残差连接
• STPE（L-STEP-inspired Spectral Temporal Positional Encoding）：rank-32 稀疏 SVD 外挂重排
• EdgeBank：记忆历史边，软融合
• CommonNeighbors：共同邻居信号，软融合

详情见 docs/pencil_fast/ 下的设计文档。

---

环境安装

• Python 版本：3.11
• 深度学习框架：Jittor 1.3.11+

pip install -r requirements.txt

注意：需要在 WSL2 Ubuntu 环境下运行，使用已配置好的 conda 环境 common（含 CUDA 12.4 + cuDNN 9.14 + Jittor 1.3.11 + JittorGeometric）。

---

数据准备

比赛数据集包含两个子集：

数据集	类型	节点数	边数	测试样本
dataset1	非二部图	23,852	751,899	61,051
dataset2	二部图	63,348 (12,708/50,640)	2,414,703	153,420

数据放置结构：

jittor_bubu/
├── dataset1/
│   ├── train.csv
│   └── test.csv
└── dataset2/
    ├── train.csv
    └── test.csv

数据文件较大，未包含在仓库中。请从计图比赛平台下载。

---

训练

Fast 版本（当前主线：PENCIL + STPE + EdgeBank + CN）

wsl
source ~/miniconda3/etc/profile.d/conda.sh
conda activate common
cd "/mnt/d/研究生/计图/jittor_bubu"

# Dataset1（启用 EdgeBank 30 天半衰期 + CommonNeighbors）
python main_pencil_fast.py --dataset dataset1 --edge_bank_half_life_days 30 --cn_enabled

# Dataset2
python main_pencil_fast.py --dataset dataset2

Pro 版本（全量训练：50 epochs / 50 negatives / temperature=0.5）

python main_pencil_pro.py --dataset dataset1
python main_pencil_pro.py --dataset dataset2

关键参数

参数	默认值	含义
--dataset	dataset1	数据集选择 (dataset1/dataset2)
--epochs	50	训练轮数
--batch_size	512	批次大小
--K1	8	子图邻居数
--L	32	序列最大长度
--d	128	Transformer 隐藏维度
--layers	3	Transformer 层数
--heads	4	注意力头数
--tau	0.12	InfoNCE 温度系数
--seed	42	随机种子

---

评测/推理

训练完成后，预测文件自动写入 dataset1/dataset1_result.csv 和 dataset2/dataset2_result.csv。

模型权重保存在 checkpoints/ 目录下。

打包提交：

# 打包两个数据集的结果
zip -j result.zip dataset1/dataset1_result.csv dataset2/dataset2_result.csv

---

结果说明

• 评测指标：MRR（Mean Reciprocal Rank），评估正样本在 100 个候选中的排名
• 最终成绩以计图比赛平台的线上评分为准
• 不同运行间因随机性可能有微小波动（±0.02 MRR），可通过 --seed 固定

---

项目结构

jittor_bubu/
├── main_pencil_fast.py          # Fast 版本训练入口（主线）
├── main_pencil_pro.py           # Pro 版本训练入口（全量）
├── pencil_fast_structural_fusion.py  # 结构融合核心模块
├── requirements.txt             # Python 依赖
├── LICENSE                      # MIT 许可证
├── .gitignore                   # Git 忽略规则
├── docs/                        # 设计文档
│   ├── pencil_fast/             # Fast 版本技术文档
│   └── pencil_pro/              # Pro 版本技术文档
├── tests/                       # 单元测试
├── archive/                     # 历史版本归档
├── 参考文献/                    # 参考论文
└── presentations/               # 演示材料

---

第三方引用

来源	说明	许可证
PENCIL	核心 Transformer 架构设计参考	MIT
L-STEP	时序位置编码（STPE）思路来源	—
Jittor / JittorGeometric	深度学习框架与图神经网络库	Apache 2.0

---

许可证

本项目采用 MIT License。