OpenGraph: 通用图大模型

论文：arxiv.org/abs/2403.01121 代码：github.com/HKUDS/OpenGraph 港大数据智能实验室主页: sites.google.com/view/chaoh

研究背景

图学习（Graph Learning）技术能够对复杂的关系数据进行挖掘和学习，在推荐系统、社交网络分析、引用网络和交通网络等多个领域都显示出了巨大的应用价值。图神经网络（Graph Neural Networks, GNNs）基于迭代的消息传递机制，能够捕捉图结构数据中的复杂高阶关系，在各类图学习应用场景中取得了巨大的成功。

通常，这种端到端的图神经网络需要大量、高质量的标注数据才能获得较好的训练效果。近年来，一些工作提出图模型的预训练-微调（Pre-training and Fine-tuning）模式，使用各种自监督学习任务在无标注的图数据上首先进行预训练，再在少量标注数据上进行微调，以对抗监督信号不足的问题。这里的自监督学习任务包括对比学习、掩码重建、局部全局互信息最大化等方法。

尽管这种预训练方法取得了一定的成功，但他们的泛化能力较为有限，特别是当预训练和下游任务出现分布偏移（Distribution Shift）时。例如在推荐系统中，预训练模型使用较早时间采集的数据，而预测时用户喜好、商品热度常常出现变化。为了解决这一问题，一些近期工作提出了针对图模型的提示微调方法，能够更加高效地将预训练模型适应到下游任务和数据上。

尽管上述工作推进了图神经模型的泛化能力，但这些模型全都假设：训练数据和测试数据有相同的节点集合和特征空间。这极大地限制了预训练图模型的应用范围。因此，本文探索进一步提升图模型泛化能力的方法。具体来说，我们希望OpenGraph能够捕捉通用的拓扑结构模式，对测试数据进行零样本预测。即仅通过前向传播过程，就可以对测试图数据进行高效的特征提取和准确预测。模型的训练过程在完全不同的图数据上进行，在训练阶段不接触测试图的任何元素，包括节点、边、特征向量。

为了达成这一目的，本文需要解决以下三个挑战：

• C1. 跨数据集的Token集合变化：零样本图预测任务的一个显著困难是，不同的图数据通常有完全不同的图token集合。具体来说，不同图的节点集合通常没有交集，并且不同图数据集也经常使用完全不同的节点特征。这使得模型不能通过学习和特定数据集的图token绑定的参数，来进行跨数据集的预测任务。
• C2. 高效的节点间关系建模：在图学习领域，节点之间常常存在错综复杂的依赖关系，模型需要对节点的局部和全局邻域关系进行综合考量。在搭建通用图模型时，一个重要的任务是能够高效地建模节点间关系，这能够增强模型在处理大量图数据时的模型效果和可扩展性。
• C3. 训练数据稀缺：由于隐私保护、数据采集成本等原因，在图学习的很多下游领域中都广泛存在数据稀缺问题，这使得通用图模型的训练容易缺乏对某些下游领域的了解而产生次优的训练效果。

为了应对上述挑战，我们提出了 OpenGraph，这是一个擅长零样本学习的模型，能够识别不同下游领域之间可迁移的拓扑结构模式。我们通过创建一个具有拓扑感知投影方案的图tokenizer来解决挑战 C1，从而生成统一的图tokens。为了应对挑战 C2，我们设计了一个可扩展的图Transformer，它配备了基于锚点采样的高效自注意力机制，并包括了token序列采样以实现更高效的训练。为了解决挑战 C3，我们利用大型语言模型进行数据增强，以丰富我们的预训练，使用提示树算法和吉布斯采样来模拟现实世界的图结构关系数据。我们在多个图数据集上进行的广泛测试显示了 OpenGraph 在各种设置中的卓越泛化能力。

模型介绍

模型整体架构如下图所示，可以分为三个部分，分别为1）统一图Tokenizer，2）可扩展的图Transformer，3）大语言模型知识蒸馏。

统一图Tokenizer

为了应对不同数据集在节点、边、特征上存在的巨大差异，我们的首要任务是建立一个统一的图tokenizer，能够有效地将不同图数据影射为统一的token序列。在我们的tokenizer中，每个token都有一个语义向量，用于描述对应节点的信息。通过采用统一的节点表征空间，以及灵活的序列数据结构，我们希望为不同的图数据进行标准化、高效的tokenization。为了达成这一目的，我们的tokenizer采用了经过平滑的拓扑信息，以及一个从节点空间到隐表征空间的映射函数。

高阶平滑邻接矩阵. 在图tokenization过程中，我们使用邻接矩阵的高次幂作为输入之一，这种方式既能够获取图结构的高阶连接关系，也能够解决原始邻接矩阵中连接稀疏性的问题。计算过程中进行了Laplacian归一化，并将不同阶的邻接矩阵幂全部考虑进来，具体计算方法如下。

对任意图的拓扑感知映射. 不同数据集的邻接矩阵在维度上存在巨大的差异，这使得我们不能直接将邻接矩阵作为输入，再使用一个固定输入维度的神经网络进行处理。我们的解决方案是首先将邻接矩阵投射为节点表征序列的形式，后续即可使用可变长的序列模型进行处理。而为了减少映射过程的信息损失，我们提出了一种拓扑感知的映射方法。

首先，我们拓扑感知映射的值域是一个维度较高的隐表征空间。之前的一些工作指出，在采用较大的隐空间维度时，即使是随机的映射也常常可以产生令人满意的表征效果。为了进一步保留图结构信息，并减少随机性影响，我们使用快速特征值分解（SVD）来构造我们的映射函数。在实际实验中，两轮的快速特征值分解可以有效地保留拓扑信息，且产生的的计算开销相对其他模块可以忽略不计。

可扩展的图Transformer

经过无参数的图tokenization过程，对具有不同特征的图数据，OpenGraph分配了统一的拓扑感知图token表征。接下来的任务是采用可训练的神经网络，对节点间的复杂依赖关系进行建模。OpenGraph采用了transformer架构，以利用其在复杂关系建模方面的强大能力。为了保证模型效率和性能，我们引入了以下两种采样技巧。

Token序列采样. 由于我们的图token序列数据一般有较大的token数量和隐表征维度，OpenGraph采用的图transformer对输入的token序列进行采样，只学习当前训练批次内的token间两两关系，使得需要建模的关系对数量从节点数量平方，降低到训练批次大小的平方，从而大大减小图transformer在训练阶段的时间和空间开销。并且，这种采样方法能够让模型在训练时更加关注当前的训练批次。尽管对输入数据进行了采样，由于我们的初始图token表征包含了节点间的拓扑关系，采样的token序列仍然能够一定程度地反映全图所有节点的信息。

自注意力中的锚点采样方法. 虽然token序列采样将复杂度从节点数量平方降低到了批次大小平方，但平方级别的复杂度对批次大小有着较大的限制，使得模型训练无法采用较大的批次，从而影响整体的训练时间和训练稳定性。为了缓解这一问题，OpenGraph的transformer部分放弃了对所有token之间两两关系的建模，而是采样部分锚点，将所有节点间的关系学习拆分成两次所有节点与锚点之间的关系学习。

大语言模型知识蒸馏

由于数据隐私等原因，获取各个领域的数据来训练通用图模型是很有挑战性的。感受到大型语言模型（LLM）所展示的惊人知识和理解能力，我们利用其力量来生成各种图结构数据，用于通用图模型的训练。我们设计的数据增强机制，使LLM增强的图数据能够更好地近似真实世界图的特征，从而提高了增强数据的相关性和有用性。

基于LLM的节点生成. 在生成图时，我们的初始步骤是创建一个适合特定应用场景的节点集。每个节点都具有一个基于文本的特征描述，该描述有助于后续的边生成过程。然而，当处理真实世界场景时，这项任务可能特别具有挑战性，因为节点集的规模很大。例如，在电子商务平台上，图数据可能包含数十亿种产品。因此，有效地使LLM生成大量节点成为一个重大挑战。

为了解决上述挑战，我们采用了一种策略，不断将一般节点分成更细粒度的子类别。例如，当生成电商场景下的产品节点时，我们首先使用类似于“列出淘宝等电子商务平台上的所有产品的子类别”的查询提示LLM。LLM回答了一个子类别列表，如“衣服”、“家居厨具”和“电子产品”等。然后，我们要求LLM进一步细化每个子类别来继续这一迭代分裂过程。这个过程一直重复，直到我们获得了类似于真实世界实例的节点，例如一个带有“衣服”、“女士衣物”、“毛线衫”、“带兜毛线衫”和“白色带兜毛线衫”标签的产品。

提示树算法. 将节点分割成子类别并生成细粒度实体的过程遵循一种树状结构。最初的一般节点（例如“产品”、“深度学习论文”）作为根，细粒度实体作为叶节点。我们采用树形提示策略来遍历和生成这些节点。

基于LLM和吉布斯采样的边生成. 为了生成边，我们使用吉布斯采样算法与上文生成的节点集。算法从一个随机样本开始进行迭代，每次在当前样本的基础上，采样对其中某一个数据维度进行改变后得到的样本。这一算法的关键在于估计在当前样本的条件下，某个数据维度改变的条件概率。我们提出根据节点生成时得到的文本特征，由LLM进行概率估计。由于边的集合空间较大，为了避免让LLM对其进行探索而产生巨大的开销，我们首先使用LLM对节点集合进行表征，再基于表征向量，使用简单的相似度算子对节点间关系进行计算。在以上的边生成框架内，我们还采用了以下三种重要的技巧进行调整。

动态概率归一化. 由于LLM表征的相似度可能与[0, 1]范围差距巨大，为了获得更适合采样的概率数值，我们使用一种动态概率归一化的方法。该方法动态维护采样过程中最近的T'个相似度估计数值，计算他们的均值和标准差，最后将当前的相似度估计映射到该均值上下两个标准差的分布范围中，从而得到近似[0, 1]的概率估计。

引入节点局部性. 我们基于LLM的边生成方法，能够有效地根据节点的语义相似性，确定他们的潜在连接关系。然而，它倾向于在所有语义相关的节点之间创建过多的连接，忽视了真实世界图中重要的局部性概念。在现实世界中，节点更有可能连接到相关节点的子集，因为它们通常只能够与一部分节点有限地互动。为了模拟这一重要特性，我们引入了一种在边生成过程中将局部性纳入考虑的方法。每个节点都随机分配一个局部性索引，两个节点之间的交互概率受到局部性索引绝对差值的衰减影响，节点的局部性索引差别越大，则衰减越严重。

注入图拓扑模式. 为了使生成的图数据更加符合拓扑结构的模式，我们在第一次图生成过程中再次生成修正的节点表征。这一节点表征使用简单的图卷积网络在初始生成图上得到，他能更好地符合图结构数据的分布特点，避免图和文本空间之间的分布偏移。最终，我们在修正的节点表征基础上，再次进行图采样，得到最终的图结构数据。

实验验证

实验中，我们仅使用基于LLM的生成数据集进行OpenGraph模型训练，而测试数据集都是各个应用场景下的真实数据集，并包括了节点分类和链路预测两类任务。实验的具体设置如下：

0-shot设置. 为了验证OpenGraph的零样本预测能力，OpenGraph在生成的训练数据集上测试，再使用完全不同的真实测试数据集进行效果测试。训练数据集和测试数据集在节点、边、特征、标注上均没有任何重合。

Few-shot设置. 由于大多数现有方法无法进行有效的零样本预测，我们采用少样本预测的方式对他们进行测试。基线方法可以在预训练数据上进行预训练，之后使用k-shot样本进行训练、微调或提示微调。

整体效果对比

在2个任务一共8个测试数据集上的测试效果如下所示。从中可以观察到：1）在跨数据集的情况下，OpenGraph的零样本预测效果相对现有方法具有较大的优势。2）在跨数据集迁移的情况下，现有的预训练方法有时甚至还不如其基础模型仅在少样本上从零开始训练，这体现了图模型获得跨数据集泛化能力的困难。

图Tokenizer研究

接下来我们探究图tokenizer设计对效果的影响。首先我们调整了邻接矩阵平滑结束，测试其对效果的影响。0阶时效果出现严重衰减，指示了采用高阶平滑的重要性。其次，我们将拓扑感知映射函数替换成其他简单方法，包括跨数据集的可学习独热id表征、随机映射、基于节点度的可学习表征。结果显示三种替代方案均效果较差，其中跨数据集学习id表征效果最差，现有工作中常用的度表征效果也衰减较大，不可学习的随机映射在所有替代方法中表现最好，但仍与我们的拓扑感知映射有很大的差距。

预训练数据集研究

为了验证基于LLM的知识蒸馏方法的有效性，我们使用不同的预训练数据集对OpenGraph进行训练，并测试其在不同测试集上的效果。本实验比较的预训练数据集包括单独去除我们生成方法中某个技巧的版本、与测试数据集不相关的两个真实数据集Yelp2018和Gowalla、以及与测试数据集相关的ML-10M数据集，从结果中可以看出：1）总体来说，我们的生成数据集能够在所有测试数据上产生较好的效果。2）所测试的三种生成技巧均起到了较为显著的提升效果。3）使用真实数据集（Yelp、Gowalla）进行训练可能反而会带来负面效果，这可能源于不同真实数据集之间的分布差异。4）ML-10M在ML-1M和ML-10M上均取得了最佳效果，这说明使用相似的训练数据集能够产生较好的效果。

Transformer中的采样技巧研究

这项实验对我们图transformer模块中的token序列采样（Seq）和锚点采样（Anc）进行了消融测试。结果显示，两种采样方法能够在训练和测试过程优化模型的空间和时间开销。在效果方面，token序列采样对模型效果有正向作用，而ddi数据集上的结果显示，锚点菜样对模型效果有负面效果。

结论

本研究的主要焦点是开发一个高度适应性的框架，能够精确地捕捉和理解各种图结构中复杂的拓扑模式。通过发挥所提出模型的潜力，我们的目的是显著提高模型的泛化能力，使其在包括各种下游应用在内的零样本图学习任务中表现出色。为了进一步提高OpenGraph的效率和健壮性，我们在可扩展的图transformer架构和基于LLM的数据增强机制的基础上构建了我们的模型。通过在多个基准数据集上进行的大量实验，我们验证了模型的杰出泛化能力。本研究在图基座模型方向作出了初步探索的尝试。在未来的工作中，我们计划赋予我们的框架自动发现噪声连接和具有反事实学习影响力的结构的能力，同时学习各种图的通用和可转移的结构模式。