中科大成果斩获图学习“世界杯”单项冠军，霸榜蛋白质功能预测任务超1年

中科大成果，拿下图学习“世界杯”单项冠军！

由中科大王杰教授团队（MIRA Lab）提出的首个具有最优性保证的大语言模型和图神经网络分离训练框架，在国际顶级图学习标准OGB（Open Graph Benchmark）挑战赛的蛋白质功能预测任务上斩获「第一名」，该纪录从2023年9月27日起保持至今。

OGB是目前公认的图学习基准数据集“标杆”，由图学习领域的国际顶级学者斯坦福大学Jure Leskovec教授团队建立，于2019年国际顶级学术会议NeurIPS上正式开源。

最近，该论文发表在人工智能顶级期刊 IEEE Transactions on Pattern Analysis and Machine Intelligence（TPAMI 2024）。

TPAMI 是目前计算机类别中影响因子最高（影响因子 20.8）的期刊之一。

该方法引入了一个十分新颖的图神经网络的逆运算，并提出标签反卷积算法来快速近似它，进而构建一个等价的损失函数，从而消除了传统语言模型和图神经网络微调方法的学习偏差。

论文和代码均放出。

引言

图广泛应用于许多重要领域，例如引文网络、商品网络和蛋白质相互作用网络。在许多实际应用中，图中的节点具有丰富且有用的属性信息。例如，引文网络中的节点（论文）、商品网络中的节点（商品）以及蛋白质相互作用网络中的节点（蛋白质）分别包含着标题/摘要、商品的文本描述和蛋白质序列等重要信息，这些信息对下游任务只管重要。而近年来兴起的许多强大的预训练模型是从这些复杂属性中捕获节点特性的重要工具之一。

为了同时编码这些属性和图结构，一个常见的架构是将预训练模型与图神经网络GNN（Graph Neural Network）串联集成在一起，其中预训练模型作为节点编码器NE（Node Encoder）对属性进行编码。如下图所示，该架构通过节点编码器将这些复杂的节点属性变成定长的低维嵌入，再将其作为节点特征输入到图神经网络以结合图结构信息。

然而，作为NE的预训练模型本身大量参数且GNN的邻居爆炸问题（neighbor explosion），两大训练难题的叠加让直接端到端联合训练NEs和GNN在实际中并不可行的。研究者们开始研究分离NEs和GNNs分离训练的范式，即先固定NEs的参数训练GNNs一定步数（GNN的训练阶段），再固定GNNs的参数训练NEs一定步数（NE的训练阶段），两步交替迭代进行。

本工作对现有的NEs和GNNs分离训练的范式进行研究，指出了现有工作在NE训练阶段，它们没有考虑GNN中的特征卷积，导致它们提出的近似损失函数与原始联合训练的目标函数并不等价，存在显著的学习偏差，进而无法收敛到最优解（详见原论文举的反例）。

为了应对这一挑战，我们提出了一种有效的标签正则化技术，即标签反卷积LD （Label Deconvolution），通过对GNN逆映射得到一种新颖的、可扩展性强的近似标签。逆映射有效地将GNN纳入NE的训练阶段以克服学习偏差，进而产生了与联合训练等效的目标函数。于是我们也进一步证明了LD收敛到了最优目标函数值，为提出的LD方法提供了理论保证。通过实验验证，LD显著优于当下最先进的方法，在国际顶级图学习标准OGB（Open Graph Benchmark）挑战赛的蛋白质功能预测任务上斩获「第一名」，该记录从2023年9月27日起保持至今。

背景介绍

大规模属性图上的节点表示学习

重点研究‍了具有丰富有用的节点属性

的图

上的节点表示学习，其中

是所有节点的集合，

 是所有边的集合。由于原始节点属性

通常是高维的文本、图像或蛋白质序列，常用的解决方法是从中提取出

维的节点特征

，如下所示：

式中

表示节点编码器（NE）的参数。由于大型的预训练模型（如：用于蛋白质序列的ESM2, 用于文本的Bert）具有强大的特征提取能力，故将其作为节点编码器f。

为了进一步编码图结构，图神经网络将节点特征

和邻接矩阵

作为如下输入：

式中：

表示

的第i行，θ表示图神经网络的参数。若

，则

，否则

。GNNs输出节点表示H。

为了简单起见，我们定义如下记号.给定一组节B，令

表示由

组成的矩阵，所有

，其中

为M的第i行。给定一个向量函数

，令

表示一个矩阵函数，其中

。

可扩展的图神经网络结合预训练节点编码器的难点

大多可扩展的图神经网络可分为基于数据采样和基于模型结构的两类思想。

图采样

为了计算节点的小批量B中的节点表示，一种常见的解决方案是对由B构造的子图进行如下采样：

其中

。

然而，现有图采样方法中使用的|G(B)|明显大于预训练NEs中使用的mini-batch的大小。如果进一步减小现有图采样方法中|B|或|G(B)|的大小来对齐mini-batch的大小，它们的性能会显著下降，如下所示。

在本实验中，预训练NEs最大batch的大小不超过12，明显小于|G(B)|。所以，通过图采样对NEs和GNNs进行联合训练是难以实现的。

从GNN中分离特征卷积

为了避免GNNs特征卷积时的内存和时间开销，一些可扩展的GNNs（例如 GAMLP 和 SAGN）首先将特征卷积从GNNs中分离出来。然后基于固定节点特征对特征卷积进行一次预处理。然而，由于节点特征是可用NEs学习的，这种想法对于NEs和GNNs的联合训练仍然是难以承受的。

预训练节点编码器结合图神经网络的常见训练范式：分离训练框架

给定节点标签Y，优化问题为

。为了避免特征卷积严重的可扩展性问题，现有的分离训练框架提出交替优化θ和β：

式中L为真实目标函数的损失函数, L'为L的近似。

GNNs的训练阶段（优化β）

当NEs的参数β固定时GNNs是可扩展的，可直接使用上述的图采样或者特征卷积分离技术来优化GNNs。

NEs的训练阶段（优化β）

现有的独立训练框架忽略了GNNs中的特征卷积来设计新的损失函数L'，例如自监督损失

或监督损失

并具有可扩展的线性层

。

值得注意的是，NEs的训练阶段不涉及公式中GNNs的参数θ。我们的方法LD和GLEM基于不同的motivation，具体来说，LD旨在恢复GNN，而GLEM旨在提高伪标签

的质量，在测试节点

上进行半监督学习。因此，我们忽略了对Y的改进，并假设LD和GLEM中的节点标签Y是相同的。

标签反卷积

针对分离训练框架中忽略GNN特征卷积的问题，我们提出了一种简单高效的标签正则化技术，即标签反卷积（Label Deconvolution，LD）。设节点标签为Y，如果任务是半监督的（指图中的部分节点标签是缺失的），即可根据预训练的NEs推理得到的固定节点特征训练GNNs。

LD是一个分离训练框架，分别对GNNs和NEs进行训练。我们将NEs的训练阶段表述为：

式中：GNN-1为GNN的逆映射。我们称

为反标签。在NEs的训练阶段参数θ是固定的，LD的核心思想是对

进行预处理，以避免在NEs的训练阶段执行多次增加内存和时间开销的操作。因此，mini-batch的训练目标为：

其中B是节点的mini-batch。

由于非线性GNN的逆映射很难精确计算，所以我们推导GNN的有效近似来替代。接下来介绍GNN的频谱公式，将线性特征卷积与GNN分离。然后通过LD参数化具有相似表达的逆标签

，进一步避免了线性特征卷积的逆映射。

基于频谱设计的GNNs

受到频谱滤波器的启发，近来出现了许多高效的GNNs架构。LD的推导也是基于频谱的GNNs，即：

其中

是执行线性特征卷积的多项式谱滤波器，

是归一化的相邻矩阵，

是非线性多层感知器。权重

要么是可学习的，要么是固定的。如[2]所示，基于频谱的GNN可以在一些温和的假设下产生任意节点预测。这些假设也适用于许多真实世界的图数据。因此训练目标变为：

上述方程保留了GNNs的可扩展的非线性变换，并预处理了图扩散矩阵

的逆矩阵。值得注意的是，在NEs的训练阶段包含了GNNs参数θ的一部分。这种结合显著减轻了NEs和GNNs联合训练的学习偏差，同时不影响可扩展性。

标签反卷积

为了进一步避免线性特征卷积的逆映射，我们提出了一个可训练的标签反卷积来生成逆标签Y(γ)。标签反卷积旨在用γ参数化Y(γ)，使得Y(γ)的表达能力类似于

，即：

这样，训练目标变为：

通过我们提出的带γ的重新参数化方法隐式地纳入了参数

。其核心思想受到Cayley-Hamilton定理的启发。下面首先介绍两个有用的引理。

引理1. 设矩阵M的特征多项式为

。若矩阵M可逆，则M的逆矩阵为：

引理2. 若矩阵MN∈Rn×n以表示成一个次数小于n的矩阵多项式，即：

由此引出命题：

命题1. 若

可逆，则

表示为

矩阵幂的线性组合，即：

因此，我们将逆标签Y(γ) 参数化为：

其中N是一个超参数，变量

是可训练参数。

直观上来看，i-hop标签

是k-hop邻居中标签的（加权）平均值。对于一个N层GNN，节点的预测（表示）不仅依赖于它的特征，而且依赖于其N跳邻居的特征。类似地，节点的特征不仅对其预测有贡献，对其N跳邻居的预测也有贡献。因此，i-hop标签可以有效缓解NEs训练阶段的学习偏差。

逆标签的小批量版本为：

式中B为节点的Mini-Batch。

下面总结了LD算法的伪代码。

NEs不同损失函数的比较

现有的分离训练框架提出了各种损失函数来逼近

，导致联合训练产生学习偏差。我们总结了在节点标签和图结构方面的学习偏差。图5展示了联合训练、LD、GIANT 和GLEM的损失函数。LD将图结构与节点标签融合生成逆标签，保持了与联合训练相似的学习行为。然而，GIANT和GLEM忽略了图结构或节点标签，导致了显著的学习偏差。

虽然LD和联合训练有着相似的学习行为，但LD在特征存储上比联合训练更高效。具体来说，为了计算mini-batch节点B数据上的损失，LD的NE以O(|B|)的内存复杂度对B中的属性进行编码。然而，联合训练的NE对大小为G|B|的采样子图中的属性进行编码，产生比LD更大的内存复杂度O(G|B|)。

下面展示了NEs训练阶段不同训练方法的复杂性以及NEs的监督信号。LD和GLEM是所有方法中速度最快、存储最高效的算法。且与GLEM相比，LD在NEs的监督信号中还考虑到了图结构。

实验结果

团队对广泛使用的开放图基准数据集OGB（Open Graph Benchmark）中的ogbn-axiv、ogbn-product和ogbn-protein进行实验，其图数据分别为引文网络、协同购买网络和蛋白质关联网络。

如下所示，LD在不同GNN backbone的三个数据集上的表现都显著优于所有的baseline。

逆标签分析

逆标签Y(γ)是真实标签与i跳邻居标签的加权和。图8绘制了微调过程中的权重γi的变化过程。逆标签往往是真实标签或者i跳邻居中i较小的标号。这是因为真实标签和i较小的i跳邻居标签仍然是所有标签中对节点分类最重要的监督信号。此外，i 较大的i跳邻居标签存在过平滑问题，即随着i的增加，i跳邻居标签可能趋于不可区分。值得注意的是，权重γi不收敛到平凡解，其中

。这说明其他跳数的标签对节点特征提取是有帮助的。

为了进一步比较逆标签和真实标签，我们在图9中展示了节点属性的相似度和标签的相似度。我们从ogbn-arxiv数据集中随机选择了几对具有高度相似文本（即文本相似度大于0.6）但标签不同（节点0和1 , 2和3 , 4和5）的节点。我们使用TF-IDF算法和余弦相似度分别来评估文本相似度和标签相似度。图4a中每对节点都具有较高的相似度，但不同对中的节点相似度较低，我们对其进行独立选择。图4b和4c表明，逆标签为具有相似文本的节点提供相似的监督信号，为具有不同文本的节点提供不同的监督信号。然而真实标签无法实现这一特性。由此可见，逆标签通过降低图结构中的标签噪声来保留真实语义属性。

作者介绍：

石志皓，2020年获得中国科学技术大学电子工程与信息科学系学士学位。现于中国科学技术大学电子工程与信息科学系的 MIRA Lab 实验室攻读博士研究生，师从王杰教授。研究兴趣包括图表示学习和AI4Science。他曾以第一作者在 TPAMI、ICLR等期刊、会议上发表论文，曾受邀在ICLR 2023做接受率约为8%的Spotlight报告。

路方华，2023年获得上海大学机械设计与自动化专业学士学位。现于中国科学技术大学电子工程与信息科学系的 MIRA Lab 实验室攻读硕士研究生，师从王杰教授。研究兴趣包括图表示学习和自然语言处理。

论文地址：https://www.computer.org/csdl/journal/tp/5555/01/10678812/20b3hKWQ3Ru

代码地址：https://github.com/MIRALab-USTC/LD

参考文献：

[1]Zhao J, Qu M, Li C, et al. Learning on large-scale text-attributed graphs via variational inference[J]. arXiv preprint arXiv:2210.14709, 2022. 

[2]Wang X, Zhang M. How powerful are spectral graph neural networks[C]//International Conference on Machine Learning. PMLR, 2022: 23341-23362.

本文来自微信公众号“量子位”，作者：MIRA Lab ，36氪经授权发布。