[論文レビュー] 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は、大規模言語モデル（LLM）とグラフニューラルネットワーク（GNN）を統合し、自己教師あり事前学習のためのマスクドグラフモデリング（MGM）とグラフインstructチューニングを用いて、テキスト属性付きグラフ（TAGs）のための統一されたクロスドメイン基盤モデルを提案する。11の多様なグラフデータセット（5ドメインにまたがる）で、専用のGNNや汎用的LLMに劣らない、またはそれ以上のSOTAまたは優れた性能を達成し、ゼロショット設定でも優れた性能を示す。

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.

研究の動機と目的

• グラフ学習モデルが通常は1つのグラフやドメインごとに訓練されるため一般化が難しいという問題に対処するため、統一された基盤モデルを開発すること。
• 自然言語を共通の表現媒体として用いて、ドメイン間で異種のグラフ特徴量およびラベル空間を統一すること。
• グラフ埋め込みと自然言語のインstructを用いたLLMのチューニングにより、未学習のグラフタスクやクラスに対するゼロショット予測を可能にすること。
• 段階的なLM-GNNアーキテクチャを用いて、大規模かつ多様なテキスト属性付きグラフ上で自己教師あり事前学習をスケーリングすること。
• ターゲットデータセットで微調整された教師ありGNNに匹敵またはそれを上回るクロスドメイン転移性能を実証すること。

提案手法

• ノード特徴量が自然言語記述であるテキスト属性付きグラフ（TAGs）を入力とし、異種の特徴空間を統一する。
• 事前学習済みLLMとGNNの段階的アーキテクチャを採用し、LLMがテキスト的ノード特徴量を処理し、GNNが構造的依存関係を捉える。
• 自己教師あり事前学習のための新規なマスクドグラフモデリング（MGM）目的関数を導入し、ノード、エッジ、または部分グラフをマスクし、文脈から予測する。
• 局所的サブグラフを抽出するためにパーソナライズドページランク（PPR）サンプリングを採用し、ドメイン間での構造的ばらつきを低減し、スケーラビリティを向上させる。
• 事前学習済みグラフ埋め込みと自然言語のインstructを連結することでグラフインstructチューニングを実施し、LLMによるゼロショット予測を可能にする。
• ユニバーサルプロンプトテンプレートを用いて、ドメインを問わずノードレベル、エッジレベル、グラフレベルの予測タスクを統一する。

実験結果

リサーチクエスチョン

• RQ1 異種の特徴量およびラベル空間を持つ多様なグラフドメインにまたがって一般化できる単一の基盤モデルは可能か？
• RQ2 大規模なテキスト属性付きグラフにおける自己教師あり事前学習に、マスクドグラフモデリング（MGM）はどの程度有効か？
• RQ3 グラフインstructチューニングは、未学習のグラフタスクやクラスに対するゼロショット予測をどの程度可能にするか？
• RQ4 段階的LM-GNNアーキテクチャは、少データ（few-shot）およびゼロショット転移設定において、標準的なGNNと比較してどのように差を示すか？
• RQ5 事前学習データの多様性およびモデル部品選択（例：LLM、GNN）が、下流性能に与える影響は何か？

主な発見

• UniGraphは、5ドメインにまたがる11のデータセットにおけるノード、エッジ、グラフレベルのタスクで、タスク特化GNNやProdygやOFAといったマルチタスク事前学習モデルを含むすべてのベースラインを上回る性能を示した。
• 111Mノードを有する ogbn-Papers100M データセットでは、微調整なしに自己教師あり表現学習でSOTA性能を達成した。
• ゼロショット転移において、Llama2 や Vicuna といったオープンソースLLMを著しく上回り、グラフ埋め込みと自然言語の間の整合性が効果的に実現されていることを示した。
• アブレーションスタディにより、MGM、潜在的損失、PPRサンプリング、GNNの各コンポonentが性能向上に寄与していることが確認され、特にPPRサンプリングがドメイン間での一般化を向上させた。
• より大きなLLM（例：DeBERTa-V3-base）およびより深いGNNは性能向上に寄与するが、推論速度の低下を伴うため、適切なハードウェアを用いればスケーラビリティは実現可能である。
• トレーニング時間およびメモリコストは、MLMを用いた標準的なLLM事前学習と同等であり、LM部が計算コストの主因であるため、大規模グラフに対しても効率的なフレームワークであることが示された。