DeepMind MAG240M-LSC 开源方案全解析：数据管线、BGRL 自监督训练与 10 折集成评估

• 人工智能
• 深度学习
• 机器学习
• 计算机视觉
• NLP
• 强化学习

导读

本文以 DeepMind 官方开源仓库中的 MAG240M-LSC 参赛方案文档 为主体，结合 config.py、experiment.py、models.py 等源码，完整梳理该方案从环境搭建、数据下载与预处理、模型训练到 10 折集成评估与生成提交文件的全部实战链路。读者读完本文后，将能够复现 DeepMind 在 OGB-LSC 学术图赛道（MAG240M-LSC）的最终测试集结果，理解其"超深 GNN + 自监督（BGRL）"的核心方法，并能根据自己的硬件条件调整配置完成训练与评估。

一、方案背景与性能概况

MAG240M-LSC 是 OGB Large-Scale Challenge（OGB-LSC）中的学术图（academic graph）赛道，任务是预测论文（paper）的学科主题。DeepMind 在 ogb_lsc/mag/ 目录下开源了其参赛方案，技术细节记载于技术报告Large-scale graph representation learning with very deep GNNs and self-supervision（arXiv:2107.09422，文中不再重复外部链接）。

从 data_utils.py 可确认数据规模常量：

• 论文（paper）节点：121,751,666
• 作者（author）节点：122,383,112
• 机构（institution）节点：25,721
• 节点总数：244,160,499
• 边总数：1,728,364,232
• 论文原始 BERT 特征维度：768
• 分类类别数：153

官方文档报告的性能与硬件指标如下（均为项目文档声明，非本文推断）：

• 最终测试集结果：由 10 折（10 folds）模型的集成（ensemble）池化得到；
• 训练：每个模型在 4x Google Cloud TPUv4 + 1x AMD EPYC 7B12 64 核 CPU @2.25GHz 上训练不超过 72 小时；
• 推理：在 4x NVIDIA V100 16GB GPU + 1x Intel Xeon Gold 6148 20 核 CPU @2.40GHz 上耗时不超过 12 小时。

二、环境搭建

官方文档建议在 fork 的deepmind_research仓库内创建 Python 虚拟环境（可能需要先安装python3-venv包），并安装全部依赖：

git clone https://gitcode.com/gh_mirrors/de/deepmind-research cd deepmind-research python3 -m venv /tmp/mag_venv source /tmp/mag_venv/bin/activate pip3 install --upgrade pip setuptools wheel pip3 install -r ogb_lsc/mag/requirements.txt

依赖清单见 requirements.txt，核心包括dm-haiku、optax、jraph（graph 网络库，固定 commit）、graph_nets、ogb、chex、ml_collections、google-cloud-storage（用于从 GCS 下载数据）等，其中jaxline、graph_nets、jraph均以 git 固定版本安装。

若使用 NVIDIA V100 GPU，还需要安装与之兼容的 jaxlib 版本：

pip install --upgrade jax jaxlib==0.1.67+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

该命令来自 README.md，用于规避 V100 上已知的 jaxlib 兼容性问题。

三、数据下载与预处理

3.1 三步数据管线

官方 README 给出的数据准备流程分三步：

第 1 步：下载原始数据集。使用 OGB 官方竞赛工具包下载 MAG240M-LSC 数据集到本地目录ROOT。

第 2 步：整理目录结构。运行 organize_data.sh，将官方下载的异构目录结构重排为扁平、命名透明的统一目录：

/bin/bash organize_data.sh -r ROOT

需要说明的是，README 中以-r ROOT示例调用，而脚本本身实际解析的参数为-i <Downloaded data dir> -o <Task root directory>（见脚本内 usage 提示），使用时可留意这一点。脚本完成后的raw目录包含 9 个文件：

• node_feat.npy（论文 768 维 BERT 特征）
• node_label.npy（论文标签）
• node_year.npy（论文年份）
• author_affiliated_with_institution_edges.npy
• author_writes_paper_edges.npy
• paper_cites_paper_edges.npy
• train_idx.npy/valid_idx.npy/test_idx.npy（训练/验证/测试索引，由 split_and_save_indices.py 切分生成）

文档将这一层数据统称为 "raw" 数据。

第 3 步：运行预处理代码。执行 run_preprocessing.sh：

/bin/bash run_preprocessing.sh -r ROOT

该脚本依次调用 4 个 Python 模块完成预处理：

1. csr_builder.py：把边表构建为 CSR 稀疏矩阵（paper_paper_b、paper_paper_b_t、author_institution、institution_author、author_paper、paper_author），供图采样与消息传递使用；
2. pca_builder.py：对论文的 768 维 BERT 特征做 PCA 降维到 129 维（在 100 万篇论文子集上估计主成分，采样 10,000 条向量，最终对全部节点按块投影），并由论文 PCA 特征重算作者、机构的 PCA 特征；
3. neighbor_builder.py：基于论文 PCA 特征构建 Annoy 近似最近邻索引（neighbor_indices.npy、neighbor_distances.npy），用于"融合"邻接关系；
4. generate_validation_splits.py：生成 10 折（k=10）交叉验证划分，输出到data/k_fold_splits。

预处理产物统一写入preprocessed目录（相关文件名定义见 data_utils.py），包括fused_node_labels.npy、fused_paper_edges.npz、merged_feat_from_paper_feat_pca_129.npy等。文档明确警告：预处理既耗时又耗内存，仅建议在需要验证完整管线时自行运行。

3.2 直接下载预处理数据与预训练权重

官方在 GCS（bucketdeepmind-ogb-lsc）上提供了预处理好的数据与最终提交的预训练权重，可通过 download_mag.py 匿名下载（无需鉴权，使用 Google Cloud Storage 匿名客户端），payload参数二选一：

# 下载预处理数据（用于训练与评估） python3 download_mag.py --task_root=${HOME}/mag --payload="data" # 下载最终提交的预训练权重 python3 download_mag.py --task_root=${HOME}/mag --payload="models"

从源码看，下载器内置了 crc32c 校验与最多 5 次重试（_MAX_DOWNLOAD_ATTEMPTS = 5，download_mag.py），并会跳过已存在的文件，断点续传友好。payload="data"下载的内容覆盖原始索引、10 折划分、融合标签、融合论文边、各类异质边以及 PCA 合并特征等全部训练所需数组。

3.3 复现最终结果

下载两个 payload 后，官方给出的复现命令是：

/bin/bash run_preprocessing.sh -r ${HOME}/mag/

即同样需要执行一次预处理管线，之后即可用预训练权重进行评估（见第六节）。

四、配置详解：读懂 config.py

训练与评估的全部超参数集中在 config.py 的get_config(debug: bool = False)中。该配置基于 Jaxline 框架（base_config+ml_collections.ConfigDict），并支持通过命令行--config.xxx=yyy覆盖任意字段。核心参数整理如下：

4.1 训练循环全局配置

4.2 数据集与在线子图采样

邻域采样的实现细节见 sub_sampler.py 与 datasets.py：按节点类型选择对应的 CSR 矩阵取邻接行，邻域规模小于阈值时直接切片，较大时使用无放回/有放回的均匀随机采样（np.random.randint对超大规模邻域避免切片），并支持去重。子图最大深度为 3（_MAX_DEPTH_IN_SUBGRAPH = 3）。

4.3 优化器与学习率

• 优化器：adamw（即 AdamW），weight_decay=1e-5，b1=0.9，b2=0.999；
• 学习率：base_learning_rate=1e-2，warmup_steps=50000，随后余弦退火至total_steps。

具体调度实现见 schedules.py：前warmup_steps步线性升温，之后按余弦衰减；评估用的 EMA 衰减率也按余弦从base_rate=0.999退火趋近于 1。

4.4 模型结构（model_config）

4.5 损失与动态批大小

训练损失为监督分类损失 + BGRL 自监督损失的组合（bgrl_loss_scale=1.0）：

loss_config=dict( bgrl_loss_config=dict( stop_gradient_for_supervised_loss=False, bgrl_loss_scale=1.0, symmetrize=True, first_graph_corruption_config=dict(feature_drop_prob=0.4, edge_drop_prob=0.2), second_graph_corruption_config=dict(feature_drop_prob=0.4, edge_drop_prob=0.2), ), ),

动态批大小按"节点数/边数/图数"三元组约束（dynamic_batch_size_config）：训练默认n_node=340*256、n_edge=720*256、n_graph=256；评估默认n_node=340*128、n_edge=720*128、n_graph=128。源码注释提示：单 GPU 显存不足时可将256降到48（对应批大小约 48，见 experiment.py 的说明）。

五、模型架构：超深 Encode-Process-Decode + BGRL 自监督

5.1 Encode-Process-Decode 主干

模型类NodePropertyEncodeProcessDecode定义在 models.py，是典型的 Encode-Process-Decode 三段式结构：

• Encode：jraph.GraphMapFeatures用两套 MLP 分别对节点、边特征做初始嵌入（_encode），训练时随即施加 dropout；
• Process：堆叠num_message_passing_steps个jraph.InteractionNetwork块（_process），每块含节点更新 MLP 与边更新 MLP（默认disable_edge_updates=True时不更新边特征），消息聚合函数由aggregation_function决定（sum/mean/softmax），并对边消息施加 dropedge（maybe_dropedge）；每个块之后以残差连接累加节点/边表示，从而支持训练非常深的 GNN；
• Decode：节点 MLP 输出 153 类的node_logits；同时通过projector与predictor两个 MLP 头输出投影（node_embedding_projections）与预测（node_projection_predictions），供 BGRL 自监督损失使用。

模型输出ModelOutput为node_embeddings / node_embedding_projections / node_projection_predictions / node_logits四元组（models.py）。

5.2 BGRL 自监督：双视角对比

BGRL（Bootstrap Your Own Latent）自监督损失实现在 losses.py：

1. 对同一子图生成两个"损坏视角"（get_corrupted_view）：按feature_drop_prob=0.4随机遮蔽特征、按edge_drop_prob=0.2把被丢弃的边重定向为填充节点的自环（losses.py）；
2. 一个视角经在线网络（online）前向，另一个视角经 EMA 目标网络（target）前向；
3. 在 L2 归一化后计算预测头与目标投影间的均方误差，symmetrize=True时两侧对称求和；
4. 自监督损失与监督交叉熵损失（node_classification_loss，带标签掩码）加权相加。

在 experiment.py 的_bgrl_loss中可以看到完整调用链：两个损坏视角分别过 online 与 EMA 网络，且只取每个子图中心节点的表示计算 BGRL 损失；监督解码则作用于原始图，并可通过stop_gradient_embedding_to_logits切断嵌入到 logits 的梯度。ratio_unlabeled_data_to_labeled_data=10.0正是为自监督提供大量无标签中心节点的关键配置。

5.3 节点特征

训练时每个子图节点的特征由 datasets.py 的sstable_to_intermediate_graph构造，包括：节点类型 one-hot（paper/author/institution，3 类）、子图深度 one-hot、论文年份、标签 one-hot（153 类）、是否为中央节点、索引与绝对索引（作者/机构偏移见 data_utils.py）。原始 768 维 BERT 特征经 PCA 降至 129 维后作为发送者特征注入（sub_sampler.py 的get_senders）。

5.4 训练循环与 EMA

训练主循环在 experiment.py 中实现：

• _update_func通过jax.value_and_grad计算梯度，jax.lax.pmean做跨设备平均，用optax的 AdamW 更新参数；
• 参数与网络状态同时维护一套EMA 副本（ema_params、ema_network_state），EMA 衰减率按余弦退火（schedules.apply_ema_decay/ema_decay_schedule）；
• 评估一律使用 EMA 参数（evaluate中取utils.get_first(self._ema_params)），这也是提升稳定性的关键技巧；
• 训练数据流采用tf.data+ 动态批处理（batching_utils.dynamically_batch），按n_node/n_edge/n_graph上限自动聚合子图。

六、训练：20 个模型与 10 折交叉验证

6.1 一键训练脚本

/bin/bash run_training.sh -r ${HOME}/mag/

该脚本为说明性脚本，官方明确声明：它只在单机运行，实际训练时需缩小批大小和/或模型规模以适配单 GPU（run_training.sh）。

脚本的核心逻辑：

• 对 k=0~9 的每一折，各训练2 个不同随机种子的模型，共 20 个模型：BASE_SEED=100，SEED_OFFSET取 0 与 10，MODEL_SEED = 100 + SEED_OFFSET + K_FOLD_INDEX；
• 每个模型在"训练集 + 该折 90% 验证集"上训练，用剩余 10% 验证数据做早停（每个折留出不同的 10%，因此每个验证节点恰好在一个模型中被留出）；
• 早停集成轮数EPOCHS_TO_ENSEMBLE=5；
• 训练命令采用--jaxline_mode="train_eval_multithreaded"，即训练与评估分线程并行，早停评估器独立运行。

6.2 官方生产环境的分布式机制

run_training.sh内的说明给出了官方提交时的分布式训练机制（文档声明）：

• 使用 4x Cloud TPU v4 实现批并行，实际批大小约为配置值的 8 倍；
• 使用独立 CPU 为每个中心节点做邻域子采样，保证 TPU 每步都有现成批次，最大化利用率；
• 在线早停评估器跑在独立的 NVIDIA V100 机器上；
• 上述整套流程对 20 个模型各复制一份；
• 在该机制下，训练约2 步/秒，50 万步在 3 天内完成。

6.3 直接运行 experiment.py

不借助 shell 脚本时，可直接以模块方式运行 experiment.py，其 docstring 提供了完整用法。公共参数可抽成变量：

DATA_ROOT=/tmp/mag/data CHECKPOINT_DIR=/tmp/checkpoint/ OUTPUT_DIR=/tmp/predictions/ K_FOLD_INDEX=0 SHARED_ARGUMENTS="--config=ogb_lsc/mag/config.py \ --config.experiment_kwargs.config.dataset_kwargs.data_root=${DATA_ROOT} \ --config.experiment_kwargs.config.dataset_kwargs.k_fold_split_id=${K_FOLD_INDEX} \ --config.checkpoint_dir=${CHECKPOINT_DIR}"

三种运行模式：

# 仅训练 python -m ogb_lsc.mag.experiment \ ${SHARED_ARGUMENTS} --jaxline_mode="train" RESTORE_PATH=${CHECKPOINT_DIR}/models/latest/step_${STEP}_${TIMESTAMP} # 训练 + 独立线程早停 python -m ogb_lsc.mag.experiment \ ${SHARED_ARGUMENTS} --jaxline_mode="train_eval_multithreaded" RESTORE_PATH=${CHECKPOINT_DIR}/models/best/step_${STEP}_${TIMESTAMP} # 用预训练模型产出预测 SPLIT="valid" # 或 "test" EPOCHS_TO_ENSEMBLE=50 # 最终提交用 50 python -m ogb_lsc.mag.experiment \ ${SHARED_ARGUMENTS} --jaxline_mode="eval" \ --config.one_off_evaluate=True \ --config.experiment_kwargs.config.num_eval_iterations_to_ensemble=${EPOCHS_TO_ENSEMBLE} \ --config.restore_path=${RESTORE_PATH} \ --config.experiment_kwargs.config.predictions_dir=${OUTPUT_DIR} \ --config.experiment_kwargs.config.eval.split=${SPLIT}

三个实用技巧（均出自 experiment.py docstring）：

1. OOM 降批：显存不足时把动态批大小降到约 48：

   SHARED_ARGUMENTS="${SHARED_ARGUMENTS} \ --config.experiment_kwargs.config.training.dynamic_batch_size_config.n_node=16320 \ --config.experiment_kwargs.config.training.dynamic_batch_size_config.n_edge=34560 \ --config.experiment_kwargs.config.training.dynamic_batch_size_config.n_graph=48"

2. 缩短启动时间：不把完整 CSR 邻接矩阵加载进内存，改用哑邻接矩阵：

   SHARED_ARGUMENTS="${SHARED_ARGUMENTS} \ --config.experiment_kwargs.config.dataset_kwargs.use_dummy_adjacencies=True"

3. 断点续训：restore_path可与train/train_eval_multithreaded模式组合，从 checkpoint 继续训练；注意早停不会计入恢复点之前的"历史最佳"表现。

此外，experiment.py 注册了信号处理器：Ctrl+C（SIGINT）保存当前模型但继续训练，Ctrl+\（SIGQUIT）保存模型并退出，便于随时留档中间模型。

6.4 快速验证：debug 模式

若只想验证代码在自己的硬件上能正确跑通，官方建议在 config.py 中设置debug=True：此时隐藏层[32]、潜变量 32、消息传递 2 层、动态批大小大幅缩小，并自动启用use_dummy_adjacencies=True，数据集迭代只取 50 个样本（datasets.py），可以快速跑通全流程。

七、评估与 10 折集成

7.1 用预训练权重评估并生成提交文件

/bin/bash run_pretrained_eval.sh -r ${HOME}/mag/

该脚本（run_pretrained_eval.sh）的执行逻辑：

• SPLIT="test"（可改为"valid"）：
  • 在test上，10 个模型分别对全部测试节点推理，最终汇总生成提交文件；
  • 在valid上，结果按 10 折"拼接"（gather），因为每个模型只评估自己被留出的那 10% 验证节点，10 个模型恰好覆盖全部验证节点；
• EPOCHS_TO_ENSEMBLE=50：每个模型对中心节点采样 50 轮不同子图做集成（单模型单轮在 test 上约 10–25 分钟，在对应验证折上约 2–3 分钟，视 GPU 而定）；
• 使用交叉验证挑选出的最佳种子BEST_SEEDS=(100 111 102 113 104 105 106 107 108 109)（每个折一个）；
• 对每个折k，以RESTORE_PATH=${MODELS_ROOT}/k${K_FOLD_INDEX}_seed${SEED}运行--jaxline_mode="eval"，把预测写入predictions_dir；
• 最后调用 ensemble_predictions.py 完成 10 折集成与提交文件生成。

7.2 集成策略：投票 + 概率平均

ensemble_predictions.py 的集成逻辑分两层：

1. 概率平均：对每个节点收集 10 个模型的 logits，softmax 后按模型轴取平均（mean_probs）；
2. 投票为主、概率平局决胜：每个模型先对各自 argmax 的类别投一票（num_votes_per_class），总分按total_score = 10 * num_votes_per_class + mean_probs计算，即票数优先、仅在同票时用平均概率破平，最后log(total_score)取 argmax 得到集成预测。

这一策略同时利用了集成的稳定性（投票）与细粒度置信度（概率）。

7.3 验证与提交

• 验证集：用ogb.lsc.MAG240MEvaluator计算官方口径的准确率（evaluate_validation）；
• 测试集：用MAG240MEvaluator.save_test_submission生成提交文件（save_test_submission_file）；
• 集成后的预测会以{split}.dill形式保存在输出目录。

评估内部的单模型推理（_evaluate_with_ensemble）在 experiment.py 中实现：对num_eval_iterations_to_ensemble轮子图采样迭代取 softmax 概率平均，得到每轮的Predictions（node_indices/labels/predictions/logits，见 losses.py），并输出top_1到top_5命中率统计。评估配置（eval.split、EMA 退火、动态批大小）对应 config.py。

八、引用

官方文档提供了一并涵盖 PCQM4M-LSC 参赛工作的 BibTeX 引用条目，可直接使用：

@article{deepmind2021ogb, author = {Ravichandra Addanki and Peter Battaglia and David Budden and Andreea Deac and Jonathan Godwin and Thomas Keck and Wai Lok Sibon Li and Alvaro Sanchez-Gonzalez and Jacklynn Stott and Shantanu Thakoor and Petar Veli\v{c}kovi\'{c}}, title = {Large-scale graph representation learning with very deep GNNs and self-supervision}, year = {2021}, journal={arXiv preprint arXiv:2107.09422}, }

九、实践要点与注意事项小结

1. 单机脚本仅用于复现流程：run_training.sh与run_pretrained_eval.sh均为单机脚本；官方训练依赖 TPU 批并行与独立 CPU 子采样，单 GPU 上务必用debug=True或显式缩小dynamic_batch_size_config（如n_node=16320, n_edge=34560, n_graph=48）验证可行性。
2. 预处理很重：完整预处理耗时耗内存，建议直接download_mag.py --payload="data"下载现成产物；自行运行预处理仅用于验证全管线。
3. 检查点机制：模型以checkpoint.dill保存于checkpoint_dir/models/{latest,best}/step_${STEP}_${TIMESTAMP}/；训练中可用 Ctrl+C 留档、Ctrl+\ 保存退出，之后用--config.restore_path续训。
4. 评估始终用 EMA 参数，EMA 衰减率随训练余弦退火，能显著提升推理稳定性。
5. 源码细节提醒：organize_data.sh 中 README 示例参数（-r）与脚本实际参数（-i/-o）不一致，且脚本内存在对未定义变量${ROOT}的引用（第 60–63 行），实际运行时应替换为${INPUT_DIR}并核对目录路径，这也是复现时最常见的踩坑点。

通过上述链路——环境搭建 → 数据下载/预处理 → 配置理解 → 模型与损失原理 → 20 模型训练 → 10 折集成评估，即可完整复现 DeepMind 在 MAG240M-LSC 上的最终提交结果，并在此基础上按需调整模型深度、批大小与集成轮数进行自己的实验。

• 人工智能
• 深度学习
• 机器学习
• 计算机视觉
• NLP
• 强化学习