专题解读 | 从图泛化(Generalization)到图大小泛化(Size Generalization)

近些年，图神经网络（GNN）由于其在图结构数据建模方面的出色表现，已经广泛应用于了各式的图相关任务，如：节点分类，链接预测和图分类等。在现实任务中，尤其是在图分类任务上，图数据的大小（图中的节点数量）往往会发生非常大的变化，这种在训练集和测试集的图大小变化往往会严重影响到GNN在测试数据上的泛化能力。我们期待设计一种在小图数据上训练也能够很好泛化到大图数据上的GNN模型，这种优秀的能力被称为图大小泛化 （Size generalization）能力。现有的一些工作基于传统的统计机器学习方法（PAC-Bayesian，VC-dimension和Rademacher complexity）已经对GNN的泛化（Generalization）能力进行了详细分析，但这些针对分析难以推广到图大小泛化的分析上，本专题将分别介绍发表在Neurips 22的《Generalization Analysis of Message Passing Neural Networks on Large Random Graphs》和ICML 23的《Wasserstein Barycenter Matching for Graph Size Generalization of Message Passing Neural Networks》这两篇论文，介绍如何从基于图泛化的分析推广到图大小泛化的分析，以及根据分析提出大小泛化能力更好的模型。

由于大部分GNN实际上可以被看作利用了类似的消息传递神经网络（MPNN）的框架，所以针对图泛化和大小泛化能力的分析将以MPNN作为图神经网络的核心框架进行。其中第一篇论文介绍的是当GNN利用MPNN作为核心框架，用Mean Pooling将节点特征转换成全图特征时，基于随机图模型生成的图数据得到的损失将满足何种Bound，而第二篇论文则考虑将非参的Mean Pooling改为参数化的Wasserstein Barycenter Matching Pooling后，当训练集和测试集的平均图大小不一样时，两者的loss bound会随图节点数增多而减小。由于这两篇论文理论分析部分较长，限于篇幅，本文只介绍两者证明的内在逻辑和核心结论，不会展开具体证明过程。

2. 论文1《Generalization Analysis of Message Passing Neural Networks on Large Random Graphs》

2.1 摘要

消息传递神经网络（MPNN）作为卷积神经网络在图结构数据上的扩展，它们迅速流行起来，现在被认为是解决各种图相关问题的最新工具。这篇文章研究了MPNN在图分类和回归中的泛化误差。其中分析假设来自不同类别的图是从不同的随机图模型（Rodom Graph Model）中采样的，我们发现，当从这种分布中采样的数据集上训练MPNN时，泛化误差会随着MPNN的复杂性增加而增加，这不仅与训练样本数量有关，还与图中节点的平均数量有关。这表明，只要图的规模足够大，具有高复杂性的MPNN可以从一个小型图数据集中泛化。泛化界限是从均匀收敛结果中导出的，该结果表明，任何应用于图上的MPNN都可以近似于应用于将图离散化的几何模型上的MPNN。

2.2 预备知识

机器学习的目的可以被看作最小化期望损失（Expected loss）：

由于我们无法得到真实的数据分布，我们只能靠从真实分布中采样得到的有限训练数据来模拟这种损失，我们称之为经验损失（Empirical loss）：

这两个损失往往会有一定的差异，这种差异被称为泛化误差（Generalization Error）：

一些文章已经针对这种泛化误差进行了分析，这种分析需要

1. 数据生成空间：假设训练数据满足某种采样分布，如：随机块模型（Stochastic Block Model），ER随机图模型（Erdős–Rényi model），随机几何图模型（Random geometric graph）
2. 模型具体细节：确定数据到预测结果的具体运算步骤
3. 确定的损失函数：确定评估基于给定数据计算得到的预测结果和真实标签差异的损失函数（分类任务的Cross Entropy，回归任务的常用的MSE）
4. Bound分析手段：用于拟合基于采样数据的经验损失和基于生成空间的期望损失之间差异界限的统计估计方法（VC-dimension based bounds，Rademacher complexity based bounds， PAC-Bayesian based bounds）

Graph

随机图模型（Random Graph Model, RGM）

以下为随机图模型的定义，可以将其理解为图数据的生成空间

离散消息传递图神经网络（Message Passing Graph Neural Network）

gMPNN作用于从图生成空间中采样出来的Graph，和常见的MPNN定义一致。

连续的消息传递神经网络（Message Passing Graph Neural Network）

cMPNN作用于整个图生成空间的模型，cMPNN可以被看作gMPNN的连续化版本，主要用于分析generalization bound。

其中，聚合函数为

2.3 理论分析

这一节将介绍根据以下假设，推导出Generalization Error：

• 数据生成空间：图数据 ，标签数据
• 模型具体细节：gMPNN（作用于Graph）和cMPNN（作用于生成空间）

推导一共分为两大部分，即

• 证明 Convergence of MPNN：采样数据通过gMPNN后的输出平均 和 数据生成空间（数据期望）通过cMPNN后的输出期望 的差异范围（difference bound）

• 证明 Generalization of MPNN : 采样数据通过gMPNN后的损失平均 和 数据生成空间（数据期望）通过cMPNN后的损失期望 的差异范围

Convergence的推导可以分为2大步（1. pooling前节点级别表示的差异，pooling后图级别表示的差异）和7小步：

• Difference bound before pooling （pooling前cMPNN和gMPNN两者在节点表征上的difference bound）
1. Simplified layer-wise difference bound （第n层节点表征的简化版difference bound）
2. Recurrence relation on layer-wise difference bound（第n-1层的bound和第n层bound的递推公式）
3. Accurate layer-wise different bound （第n层节点表征的完整版difference bound）
4. Eliminate undesired constant （常数项化简，去掉bound中不想要的常数项）
• Difference bound after Pooling（pooling后的cMPNN和gMPNN两者在图表征上的difference bound） 5. Simplified pooling bound with a probability (简化版的difference bound) 6. Pooling bound with a probability （完整版difference bound） 7. Revision for the rest situation （所有可能情况下的平均difference bound）

相关详细证明见《Generalization Analysis of Message Passing Neural Networks on Large Random Graphs》原文

限于篇幅，我们给出每一大步的推导得出的结果

Pooling 前节点级别表示差异范围

其中和分别表示具有相同参数的gMPNN和cMPNN，表示作用在图生成空间上的信号函数（可以理解成生成节点的特征的函数），表示采样图的平均图节点数，表示图生成空间的Minkowski维度（表示图生成空间的复杂程度），也表示是图生成空间的相关常数， 和 分别表示图信号函数在图生成度量空间上Lipschitz常数，以及其无穷范数。则由消息传递神经网络的相关常数（消息生成函数和节点表示更新函数的偏置，lipschitz常数等）组成。

容易发现，在pooling前节点级别表示平均（empirical）与节点表示期望（expected）的差异主要被给bound住了（直观的理解是：当足够大时，的增长率不足以抵消项的衰减），这表明当训练数据中的图节点数增加时，这种节点表征层面的差异会减小，我们推测这种情况在图级别表征上也会存在，原因在于：当进行Mean pooling时，更多参与pooling的节点能够让最终得到的图表示平均（empirical）更好地拟合图表示期望（expected）。

Pooling 后图级别表示差异范围

其中和分别表示带有Mean pooling层的gMPNN和cMPNN，，和同样由消息传递神经网络的相关常数组成。

和上面分析类似，我们能够发现在pooling后图表示平均（empirical）与图表示期望（expected）的差异主要也被给bound住了。

在这一步的推导分析中，我们在保留上述假设的同时，额外假设采样图数据的标签分布满足多项式分布。

结合多项式分布的集中不等式

我们建立起上面推倒的图级别表示差异与任务相关的损失差异的关系。

最终推导结果如下：

同样，容易发现采样得到的损失平均和损失期望的差异同样也会被限制住（定理给出的结果是差异的平方，真实结果需要开方）。

2.4 总结

到此我们建立起了当神经网络为MPNN，采用Mean Pooling策略得到图表示时，generalization error和图大小的关系。

3. 论文2《Wasserstein Barycenter Matching for Graph Size Generalization of Message Passing Neural Networks》

3.1 摘要

图大小的泛化对于消息传递神经网络（MPNNs）来说是一项具有挑战性的任务。MPNNs 在不同的图大小上的图级别分类性能会下降。最近的理论研究表明，与图大小相关的慢速不可控的收敛速度可能会对大小泛化产生不利影响。为了解决由底层维度信号生成空间中节点之间的相关性导致的不可控收敛速度，我们提出使用Wasserstein重心作为图级别的共识来应对节点级别的相关性。方法上，我们提出了一种称为Wasserstein重心匹配（WBM）层，它通过计算输入图的MPNN过滤的节点嵌入与一些经过学习的按类别划分的重心之间的Wasserstein距离来表示输入图。从理论上讲，我们证明了具有WBM层的MPNN的收敛速度是可控的，并且与信号生成空间的维度无关。因此，具有WBM层的MPNN对于缓慢不可控的收敛速度和大小变化更加具有鲁棒性。从经验上看，WBM层显著改善了在真实世界图数据集上使用不同基础架构（例如GCN、GIN和PNA）的普通MPNNs的大小泛化性能。

3.2 与论文1联系

• 论文1推导出了在采用Mean pooling时，MPNN的Generalization Error会被给bound住，由于取决于图生成空间的性质，这会使得这种Generalization Bound变得不可控。
• 论文2则在此基础上完成了
• 通过三角不等式，建立起了Generalization Error和Size Generalization Error的联系
• 提出WBM pooling层替代Mean pooling层，使得模型的Size Generalization Error能被更可控的和给bound住（其中为pooling层之前节点表示维度，能够作为超参数被控制）
• 完成在采用WBM pooling后，MPNN的Generalization Error和Size Generalization的推导

3.3 预备知识

Wasserstein 距离可以被看作不同分布距离的一种度量。相较于散度和等其他分布距离度量方式，它具有如下优势：

• 度量时能够考虑不同分布分布之间的几何特性
• 度量时能够更好的刻画相距较远分布的距离

其具体定义如下：

Wasserstein 重心可以被看作一系列分布的中心分布，其满足到所有分布的Wasserstein距离之和最小的特点，可以类比于欧氏中心是到所有点的欧式距离之和最小的点的定义。其具体定义如下

详细中文科普请参考Wasserstein Distance

3.4 模型

其中为类别为的图所属的生成（测度）空间，为类别为的图生成（测度）空间的Wasserstein重心估计。

另外我们不会直接根据barycenter公式计算，而是采用了计算上更高效的梯度下降求解，loss的公式化表达如下：

WBM pooling层的公式化表达为：

3.5 理论分析

类似于论文1中的推导过程，论文2同样首先推导了新架构下MPNN的Convergence，再通过多项式分布的集中不等式建立起了Convergence和Generalizaiton Error的联系。

不同于论文1中从图生成角度开始建模，先推导出节点表示差异后，再推导Mean pooling后的图表示差异。论文2则直接对节点表示的分布进行假设，即假设通过MPNN后得到的节点表示满足亚高斯分布。这种对节点表示分布的假设能够帮助后续推导时直接使用于Wasserstein distance和barycenter相关的不等式。限于篇幅，最终推导结果如下：

类似于论文1中的推导，本文通过了多项式分布的集中不等式建立起了convergence和generalization的关联，最终推导结果如下：

容易看出Generalization Error能够有效的被中的（低维）和（高维）给bound住。

不同于Generalization中建模的是采样数据到真实数据的差异，size generalization则是在假定训练数据的图大小于测试数据的图大小不同时，具有相同参数的模型在训练数据的loss和在测试数据上的loss的差异。

利用三角不等式建立起两者联系

其中 和 分别表示采样得到的训练数据和测试数据，两者中的图平均节点数分别为和。

3.6 实验

相比于采用Mean pooling的方案，改用WBM pooling后，不同的Backbone在graph size shift的任务场景下均有明显表现提升

相比其他针对size generalization进行模型改进的工作，WBM pooling layer的改进也具有明显的优势。

容易发现，模型呈现随的维度增加，其在测试集上的表现先上升后下降的趋势。这主要是因为当过小时，模型参数量过少，难以学习到准确的数据到标签的映射关系，而当过大时，size generalization error bound也会随之变大，这使得模型size generalization能力下降，导致同样会在测试集上表现下降。

4. 总结

到此我们通过这两篇文章探究了Generalization和Size Generalization，以及他们之间的关联，并成功通过替换Mean pooling为WBM pooling，将不可控的size generalization error bound变成了可控的error bound，这使得模型的size generalization能力变得可控，最后通过实验验证了该方法的有效性。