Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs¶

会议: NeurIPS 2026
arXiv: 2610.10520
OpenReview: https://openreview.net/forum?id=g12F7YzATu
领域: 模型压缩
关键词: GNN-to-MLP蒸馏、图几何、Ollivier-Ricci曲率、Wasserstein对齐、谱失真

一句话总结¶

针对传统GNN到无图MLP蒸馏忽略拓扑几何导致的稀疏图谱欠拟合与稠密图谱过拟合问题，本文提出基于Ollivier-Ricci曲率自适应加权的Wasserstein几何蒸馏框架G2MLP，在保持纯MLP高速推理的同时显著超越图教师与基线蒸馏模型。

研究背景与动机¶

图神经网络（GNNs）依赖邻居消息传递机制聚合结构信息，在各领域取得了强大表征能力，但在欺诈检测、高并发推荐等严苛低延迟场景下，多跳依赖与邻居递归抓取带来了巨大的推理开销与内存瓶颈。为消除推理时对图拓扑的依赖，以GLNN为代表的GNN-to-MLP知识蒸馏技术应运而生：在训练期借助教师GNN监督结构无关的MLP学生，推理期仅需将节点特征送入MLP前向传播，从而实现数十倍甚至上百倍的加速。随后的一系列工作尝试注入位置编码、可学习结构码本或基于置信度的样本重加权，但它们几乎无一例外地假设图结构具有全局均匀性，对全图所有节点施加同质的蒸馏监督。

然而，消息传递在不同图区域诱导的几何变形存在本质异质性。处于图拓扑边界的节点蕴含对判别至关重要且难以仅凭自身特征推断的高频图引导方向，而处于高连通内部稠密子图的节点则在过度聚合下被强烈平滑并向类别原型收缩。这种拓扑异质性导致传统基于逐节点KL散度对齐的MLP学生在特征空间出现双重谱失真：在低曲率的稀疏图或边界区域，学生无法捕捉教师模型中由拓扑诱导的高能量有效方差，出现“谱欠拟合”（有效秩亏缺）；而在高曲率的内部稠密区域，教师表征已平滑坍缩为低秩流形，但仅凭特征输入的学生却仍保留着未被图正则约束的虚假杂散方差，出现“谱过拟合”（有效秩虚高）。由于逐点KL散度仅对齐单节点边缘预测分布，无法约束隐藏空间的邻域表征分布，无法抑制上述谱失真。

本文的切入角度是：利用神经正切核（NTK）理论导出能量加权的理想几何对齐目标，证明节点对齐监督应根据其谱误差富集的拓扑位置进行差异化分配。核心 idea：以局部可计算的Ollivier-Ricci曲率（ORC）作为表征失真代理指标，构建曲率自适应Wasserstein对齐目标（G2MLP），对低曲率边界节点强化预测级监督以弥合谱欠拟合，对高曲率内部节点强化对角高斯邻域表征分布对齐以抑制谱过拟合，在保持部署期零图依赖的前提下实现精准几何迁移。

方法详解¶

整体框架¶

G2MLP包含两阶段训练与纯MLP部署流程。在预处理与教师训练阶段，预先离线计算并缓存全图节点的Ollivier-Ricci曲率指标，同时训练标准GNN教师（如GraphSAGE或Graph Transformer），缓存教师输出的隐藏表征与软标签。在学生蒸馏阶段，学生MLP仅以原始节点属性为输入，前向输出隐藏嵌入与预测分类Logits。训练损失通过曲率自适应权重动态划分：通过负曲率温度将预测级对齐集中分配给拓扑边界节点，通过正曲率温度将表征级邻域Wasserstein对齐集中分配给拓扑稠密内部节点。推理期彻底抛弃图拓扑结构，仅运行标准MLP前向传播。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["图拓扑与节点属性<br/>输入原始图数据"] --> B["阶段1：离线拓扑曲率计算与教师预训练<br/>解算ORC与冻结GNN教师表征"]
    B --> C["阶段2：学生前向传播与局部测度建模<br/>纯MLP前向构建局部邻域分布"]
    C --> D["曲率自适应分配机制<br/>按曲率分流预测级与表征级监督权重"]
    D --> E["停止梯度对齐目标优化<br/>正则化CAW损失更新MLP参数"]
    E --> F["部署期纯MLP推理<br/>零图访问与58倍推理加速"]

关键设计¶

1. 局部邻域测度的二阶矩对角高斯对齐：高效低阶分布匹配

针对逐点KL散度无法捕捉流形几何且离散Wasserstein距离计算复杂度高达立方级的问题，本文为每个节点构建基于一跳闭邻域的局部经验表示测度 \(\widetilde{\nu}_{v}\)，并将其松弛为对角协方差的高斯分布进行二阶矩匹配。对于节点 \(v\) 及其闭邻域 \(\mathcal{N}^{+}(v)\)，先将节点隐藏嵌入 \(\mathbf{h}_u\) 与前分类Logits \(\mathbf{z}_u\) 拼接为复合向量 \(\mathbf{r}_u = [\mathbf{h}_u^\top, \mathbf{z}_u^\top]^\top\)，并在邻域内计算均值向量 \(\mathbf{m}_v\) 与逐坐标方差向量 \(\boldsymbol{\sigma}_v\)：

\[\mathbf{m}_v = \frac{1}{|\mathcal{N}^{+}(v)|} \sum_{u \in \mathcal{N}^{+}(v)} \mathbf{r}_u, \quad \boldsymbol{\sigma}_v = \frac{1}{|\mathcal{N}^{+}(v)|} \sum_{u \in \mathcal{N}^{+}(v)} (\mathbf{r}_u - \mathbf{m}_v)^2\]

根据对角高斯分布在二阶Wasserstein度量下的显式闭式解，教师与学生局部邻域分布之间的距离简化为均值欧氏距离与标准差残差之和：

\[\mathcal{W}_{2,\mathrm{diag}}^{2}(\widetilde{\nu}_{v}^{T}, \widetilde{\nu}_{v}^{S}) = \|\mathbf{m}_{v}^{T} - \mathbf{m}_{v}^{S}\|_{2}^{2} + \|\sqrt{\boldsymbol{\sigma}_{v}^{T} + \epsilon} - \sqrt{\boldsymbol{\sigma}_{v}^{S} + \epsilon}\|_{2}^{2}\]

由于去除了非对角协方差交叉项，单节点的对齐计算复杂度从传统离散最优传输的 \(\mathcal{O}(|\mathcal{N}^{+}(v)|^3)\) 骤降至 \(\mathcal{O}(d)\)，全图单轮迭代仅需 \(\mathcal{O}((|\mathcal{V}| + |\mathcal{E}|) \cdot d)\)，使邻域几何匹配扩展到十万级节点大图成为可能。

2. 拓扑曲率符号对立自适应加权机制：针对性纠正双重谱失真

针对统一均匀加权会导致边界判别信号与内部平滑信号互相干扰的问题，本文引入Ollivier-Ricci曲率（ORC）标定局域几何特征，并分别对Logits对齐与隐层表征对齐设计了符号对立的指数加权函数。根据理论分析，ORC通过惰性随机游走测度间的 \(W_1\) 距离定义 \(\kappa(i, j) = 1 - W_1(\mu_i, \mu_j)\)；节点平均曲率标准化为 \(\widetilde{\kappa}_v\)。在低曲率的拓扑边界（\(\widetilde{\kappa}_v < 0\)），邻居差异大、教师蕴含关键的高频判别边界，因此预测级权重 \(\alpha_v\) 呈负指数放大；在高曲率的稠密核心区（\(\widetilde{\kappa}_v > 0\)），邻居高度内聚但特征易过平滑，因此表征级权重 \(\beta_v\) 呈正指数放大：

\[\alpha_{v} = \frac{\exp(-\gamma_{z} \widetilde{\kappa}_{v})}{\frac{1}{|\mathcal{V}|} \sum_{u} \exp(-\gamma_{z} \widetilde{\kappa}_{u})}, \qquad \beta_{v} = \frac{\exp(+\gamma_{h} \widetilde{\kappa}_{v})}{\frac{1}{|\mathcal{V}|} \sum_{u} \exp(+\gamma_{h} \widetilde{\kappa}_{u})}\]

其中超参数 \(\gamma_z, \gamma_h > 0\) 分别控制预测端与表征端的分流锐度，且均值归一化保证全局损失尺度恒定。由此导出的曲率自适应Wasserstein目标（CAW）形式为：

\[\mathcal{L}_{\mathrm{CAW}} = \frac{1}{|\mathcal{V}|} \sum_{v \in \mathcal{V}} \left[ \beta_{v} \mathcal{W}_{2,\mathrm{diag}}^{2}\left(\widetilde{\nu}_{h,v}^{T}, \widetilde{\nu}_{h,v}^{S}\right) + \alpha_{v} \mathcal{W}_{2,\mathrm{diag}}^{2}\left(\widetilde{\nu}_{z,v}^{T}, \widetilde{\nu}_{z,v}^{S}\right) \right]\]

该机制在拓扑边界压制谱欠拟合，在稠密内部抑制谱过拟合，使学生自适应对齐教师的有效能量流形。

3. 梯度截断动态归一化训练准则：跨尺度平稳联合优化

为避免特征范数在不同数据集间剧烈波动而造成Wasserstein损失梯度反客为主、破坏概率单纯形上的校准，本文引入带梯度截断（stop-gradient）的自适应损失缩放项。总优化目标由常规软标签KL散度损失与归一化后的曲率自适应Wasserstein损失复合而成：

\[\mathcal{L} = \lambda_{\mathrm{CE}} \mathcal{L}_{\mathrm{CE}} + \lambda_{\mathrm{KD}} \mathcal{L}_{\mathrm{KD}} + \lambda_{\mathrm{CAW}} \frac{\mathcal{L}_{\mathrm{CAW}}}{\operatorname{sg}(\mathcal{L}_{\mathrm{CAW}}) + \epsilon_{\mathrm{sg}}}\]

其中 \(\operatorname{sg}(\cdot)\) 表示停止梯度回传操作，\(\epsilon_{\mathrm{sg}} = 10^{-8}\) 防止除零异常。在节点分类与图Transformer蒸馏任务中固定 \(\lambda_{\mathrm{CE}}=0\)（纯蒸馏协议），该设计使 \(\mathcal{L}_{\mathrm{CAW}}\) 的有效梯度模长恒定归一化至批次尺度常数，既锁定了曲率加权梯度指引的方向，又解耦了超参数 \(\lambda_{\mathrm{CAW}}\) 对网络隐层维度和初始特征幅值的敏感性。

损失函数 / 训练策略¶

模型采用两阶段优化：第一阶段离线以交叉熵训练教师GNN至收敛，并使用GraphRicciCurvature库解算节点ORC指标并落盘持久化；第二阶段冻结教师参数，使用Adam优化器优化学生MLP。各基准数据集的学生学习率、权重衰减和Dropout完全沿用GLNN开源配置；\(\lambda_{\mathrm{KD}}\) 固定为 1.0，\(\gamma_h = \gamma_z = 1.0\)，仅在 \(\{10^{-2}, 10^{-1}, 5 \times 10^{-1}\}\) 内微调 \(\lambda_{\mathrm{CAW}}\)。训练实施早停策略（验证集准确率50轮无提升即停止），推理时仅保留学生前向网络。

实验关键数据¶

主实验¶

在六个主流基准节点分类数据集上，以2层（Arxiv为3层）GraphSAGE为教师，评估无图学生模型的直推（Transductive）测试准确率（%）。基准对比涵盖原始MLP、经典GLNN、可靠性加权KRD以及频域分离蒸馏FF-G2M：

数据集	指标	SAGE (教师)	GLNN	KRD	FF-G2M	G2MLP (本文)
Citeseer	Acc (%)	70.49 ± 1.53	71.22 ± 1.50	72.84 ± 1.70	72.85 ± 1.59	74.51 ± 2.35
Pubmed	Acc (%)	75.56 ± 2.06	75.59 ± 2.46	77.01 ± 3.11	76.56 ± 3.41	78.17 ± 2.75
Cora	Acc (%)	80.64 ± 1.57	80.26 ± 1.66	82.27 ± 1.31	82.38 ± 1.41	82.54 ± 1.87
A-computer	Acc (%)	82.82 ± 1.37	82.71 ± 1.18	82.87 ± 0.87	83.67 ± 1.04	83.87 ± 1.09
A-photo	Acc (%)	90.85 ± 0.87	91.95 ± 1.04	91.95 ± 1.44	92.51 ± 0.56	92.85 ± 0.78
ogbn-arxiv	Acc (%)	70.73 ± 0.35	63.75 ± 0.48	59.26 ± 0.51	58.51 ± 0.35	64.57 ± 0.31

消融实验¶

为验证各模块的实际增益，作者在三个引文网络上开展了消融测试，评估均匀Wasserstein匹配以及单侧曲率加权的作用：

配置	Cora	Citeseer	Pubmed	说明
GLNN baseline	80.26 ± 1.66	71.22 ± 1.50	75.59 ± 2.46	纯软标签KL散度蒸馏基线
+ \(W_2\) uniform (无曲率加权)	80.86 ± 2.30	71.07 ± 2.19	75.22 ± 2.15	全局无差异增加二阶高斯Wasserstein约束，收益微弱甚至略有回落
++ 曲率加权 (仅Logits端 \(\alpha_v\))	81.78 ± 1.55	73.87 ± 2.09	77.27 ± 2.92	仅对边界节点强化Logits监督，纠正谱欠拟合，性能显著跃升
++ 曲率加权 (仅特征端 \(\beta_v\))	81.34 ± 1.88	73.99 ± 2.02	77.14 ± 2.93	仅对内部节点强化表征邻域对齐，抑制谱过拟合，带来同等幅度提升
Full G2MLP (双侧曲率协同)	82.54 ± 1.87	74.51 ± 2.35	78.17 ± 2.75	双侧相反符号协同加权，较GLNN大幅提高1.6~3.3个百分点

关键发现¶

无图学生全面超越图教师：在除大规模时序图ogbn-arxiv外的5个数据集中，G2MLP在无任何图拓扑接入的情况下，准确率全面反超GraphSAGE教师（例如Citeseer上超出4.02 pp，Pubmed超出2.61 pp），证明消除过度平滑噪声的几何蒸馏具有类似正则化提纯的效果。
归纳泛化优势尤为突出：在Inductive归纳评测中，G2MLP在未见测试节点上的增益比直推模式更显著（例如在Cora/Citeseer/Pubmed/A-computer上较强基准取得2.07~2.69 pp的飞跃），表明曲率导向学到的邻域几何分布比记忆逐点软标签更具泛化鲁棒性。
双侧曲率对齐高度互补：消融实验清晰证实，缺乏曲率引导的简单高斯矩匹配几乎无增益；而单独开启边界预测对齐（\(\alpha_v\)）或内部表征平滑（\(\beta_v\)）均能带来1.5~2.8 pp的提点，两者双向配合达到最佳效果。
极端提速与零部署代价：在17万节点的ogbn-arxiv上，部署纯MLP的推理延迟由GraphSAGE的499.7 ms降低至8.6 ms，达成58倍实时加速比。

亮点与洞察¶

从NTK理论揭示GNN蒸馏的双重谱失真：首次通过神经正切核在拓扑局域的不对称性，严谨定义了稀疏图上的“谱欠拟合”（丢失高能边界模式）与稠密图上的“谱过拟合”（保留已被平滑坍缩的杂散模式），为启发式蒸馏策略奠定了理论根基。
曲率与表征粒度的精妙反向绑定：巧妙利用Ollivier-Ricci曲率的几何意义，设计了符号相反的双向指数加权，使“边界强抓Logits判别、内部强抓隐层分布凝聚”成为数学上闭环且易于计算的优雅实现。
零额外推理代价与架构通用性：所有图拓扑解算与邻域测度计算全部收敛于训练期，模型部署无需修改任何MLP架构或增加输入维度；此外该范式天然平替推广至Graph Transformer教师与边级别的链路预测任务。

局限与展望¶

单峰高斯假设的分布拟合偏差：当前邻域特征Wasserstein距离依赖单峰对角高斯矩匹配假设，在高度异质图（Heterophilic Graphs）或多模态属性分布的邻域中可能发生高阶矩丢失，导致拟合精度下降。
大规模全图精确ORC计算复杂度：精确计算全图边的Ollivier-Ricci曲率依赖线性规划求解离散最优传输，时间复杂度与平均度数的三次方成正比。虽然可离线预处理，但在超大规模图上仍需依赖Sinkhorn逼近或转用组合学Forman-Ricci曲率作为轻量替代方案。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从谱分析视角定性阐述GNN-to-MLP蒸馏的双重失真，并引入微分几何ORC曲率实现跨粒度对齐。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖节点分类直推与归纳设定、大图扩展、GT教师蒸馏、链路预测扩展以及详尽的曲率/双侧消融分析。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从NTK分析、能量对齐目标到曲率代理的转化逻辑严密流畅。
价值: ⭐⭐⭐⭐⭐ 彻底兼顾了极致的推理性能（纯MLP 58倍加速）与强大的几何表征能力，对工业界图算法落地具有高实用价值。