[论文解读] UniGraph: Learning a Unified Cross-Domain Foundation Model for Text-Attributed Graphs

Yufei He, Sui, Yuan+2 more|arXiv (Cornell University)|Feb 21, 2024

Advanced Graph Neural Networks被引用 4

一句话总结

UniGraph 通过结合大型语言模型（LLMs）和图神经网络（GNNs），并采用掩码图建模（MGM）进行自监督预训练，以及通过图指令微调实现零样本迁移，提出了一种用于文本属性图（TAGs）的统一跨域基础模型。它在涵盖5个领域的11个多样化图数据集上实现了最先进或更优的性能，包括一个包含1.11亿个节点的图，在零样本设置下优于专用GNNs和通用LLMs。

ABSTRACT

Foundation models like ChatGPT and GPT-4 have revolutionized artificial intelligence, exhibiting remarkable abilities to generalize across a wide array of tasks and applications beyond their initial training objectives. However, graph learning has predominantly focused on single-graph models, tailored to specific tasks or datasets, lacking the ability to transfer learned knowledge to different domains. This limitation stems from the inherent complexity and diversity of graph structures, along with the different feature and label spaces specific to graph data. In this paper, we recognize text as an effective unifying medium and employ Text-Attributed Graphs (TAGs) to leverage this potential. We present our UniGraph framework, designed to learn a foundation model for TAGs, which is capable of generalizing to unseen graphs and tasks across diverse domains. Unlike single-graph models that use pre-computed node features of varying dimensions as input, our approach leverages textual features for unifying node representations, even for graphs such as molecular graphs that do not naturally have textual features. We propose a novel cascaded architecture of Language Models (LMs) and Graph Neural Networks (GNNs) as backbone networks. Additionally, we propose the first pre-training algorithm specifically designed for large-scale self-supervised learning on TAGs, based on Masked Graph Modeling. We introduce graph instruction tuning using Large Language Models (LLMs) to enable zero-shot prediction ability. Our comprehensive experiments across various graph learning tasks and domains demonstrate the model's effectiveness in self-supervised representation learning on unseen graphs, few-shot in-context transfer, and zero-shot transfer, even surpassing or matching the performance of GNNs that have undergone supervised training on target datasets.

研究动机与目标

• 为解决图学习模型通常针对单个图或领域进行训练而缺乏泛化能力的问题，开发一种统一的基础模型。
• 利用自然语言作为通用表示媒介，统一不同领域之间的异构图特征与标签空间。
• 通过使用图嵌入微调大型语言模型（LLMs）的指令，实现在未见图任务和类别上的零样本预测。
• 通过级联的LM-GNN架构，在大规模、多样化的文本属性图上实现自监督预训练的扩展。
• 展示超越或匹配在目标数据集上微调的监督GNNs的跨域迁移性能。

提出的方法

• 以文本属性图（TAGs）作为输入，其中节点特征为自然语言描述，以统一不同领域之间的异构特征空间。
• 采用预训练LM与GNN的级联架构，LM处理文本节点特征，GNN捕捉结构依赖关系。
• 提出一种新颖的掩码图建模（MGM）目标，用于自监督预训练，通过掩码节点、边或子图并从上下文中预测它们。
• 使用个性化PageRank（PPR）采样提取局部子图，降低不同领域间的结构方差，提升可扩展性。
• 通过将预训练的图嵌入与自然语言指令拼接，进行图指令微调，使LLMs能够通过提示实现零样本预测。
• 使用通用提示模板，统一不同领域中的节点级、边级和图级预测任务。

实验结果

研究问题

• RQ1 一个单一的基础模型能否在具有异构特征空间和标签空间的多样化图领域中实现泛化？
• RQ2 掩码图建模（MGM）在大规模文本属性图上的自监督预训练中有多有效？
• RQ3 图指令微调在多大程度上能够实现对未见图任务和类别的零样本预测？
• RQ4 在少样本和零样本迁移设置中，级联LM-GNN架构与标准GNN相比表现如何？
• RQ5 预训练数据多样性以及模型组件选择（如LMs、GNNs）对下游性能有何影响？

主要发现

• UniGraph 在来自5个领域的11个数据集的节点、边和图级任务上均优于所有基线模型，包括任务特定的GNNs和多任务预训练模型（如Prodigy和OFA）。
• 在 ogbn-Papers100M 数据集（1.11亿个节点）上，UniGraph 在无需任何微调的情况下实现了自监督表示学习的最先进性能。
• 在零样本迁移中，UniGraph 显著优于开源LLMs（如Llama2和Vicuna），表明图嵌入与自然语言之间实现了有效对齐。
• 消融实验确认，所有组件——MGM、潜在损失、PPR采样和GNN——均对性能有显著贡献，其中PPR采样提升了跨领域的泛化能力。
• 更大的LM（如DeBERTa-V3-base）和更深的GNN可提升性能，尽管会带来推理速度的权衡，表明在适当硬件下具备可扩展性。
• 训练时间和内存成本与标准LLM预训练（MLM）相当，表明LM组件主导计算成本，使该框架在大规模图上具有高效性。