[논문 리뷰] UniGraph: Learning a Unified Cross-Domain Foundation Model for Text-Attributed Graphs

Yufei He, Sui, Yuan외 2명|arXiv (Cornell University)|2024. 02. 21.

Advanced Graph Neural Networks인용 수 4

한 줄 요약

UniGraph는 대규모 언어 모델(LM)과 그래프 신경망(GNN)을 결합하고, 자기지도 사전 훈련을 위한 마스킹 그래프 모델링(MGM)과 영향 없는 전이를 위한 그래프 지시 훈련을 통해 텍스트 속성 그래프(TAGs)를 위한 통합된 다영역 기초 모델을 제안한다. 이는 5개 영역에 걸쳐 11개의 다양한 그래프 데이터셋에서 최신 기술 수준 또는 그 이상의 성능을 달성하며, 1,110만 개 노드를 포함한 그래프에서도 전문화된 GNN과 일반 목적의 LLM보다 영향 없는 설정에서 뛰어난 성능을 보인다.

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.

연구 동기 및 목표

• 일반적으로 단일 그래프 또는 영역별로 훈련되는 그래프 학습 모델의 일반화 부족 문제를 해결하기 위해 통합된 기초 모델을 개발한다.
• 자연어를 공통 표현 매체로 사용하여 다양한 영역 간의 이질적인 그래프 특징 및 레이블 공간을 통합한다.
• 그래프 임베딩과 함께 자연어 지시어를 사용하여 LLM의 지시 훈련을 통해 새로운 그래프 작업과 클래스에 대한 영향 없는 예측을 가능하게 한다.
• 계층적 LM-GNN 아키텍처를 사용하여 대규모이고 다양한 텍스트 속성 그래프에서 자기지도 사전 훈련을 스케일링한다.
• 목표 데이터셋에서 미세조정된 감독 기반 GNN과 비교하여 초과하거나 동등한 다영역 전이 성능을 입증한다.

제안 방법

• 노드 특징이 자연어 기술서인 텍스트 속성 그래프(TAGs)를 입력으로 사용하여 다양한 영역 간의 이질적인 특징 공간을 통합한다.
• 사전 훈련된 LM과 GNN의 계층적 아키텍처를 사용하며, LM은 텍스트 노드 특징을 처리하고 GNN은 구조적 의존성을 캡처한다.
• 자기지도 사전 훈련을 위한 새로운 마스킹 그래프 모델링(MGM) 목적함수를 도입하여 노드, 간선 또는 부분그래프를 마스킹하고 맥락에서 복원한다.
• 개인화된 페이지랭크(PPR) 샘플링을 사용하여 국소 부분그래프를 추출하여 영역 간의 구조적 변동성을 줄이고 확장성을 향상시킨다.
• 사전 훈련된 그래프 임베딩과 자연어 지시어를 연결하여 그래프 지시 훈련을 수행함으로써 LLM을 통한 영향 없는 예측을 가능하게 한다.
• 일관된 프롬프트 템플릿을 사용하여 다양한 영역 간의 노드 수준, 간선 수준, 그래프 수준의 예측 작업을 통합한다.

실험 결과

연구 질문

• RQ1 다양한 그래프 영역 간에 이질적인 특징 공간과 레이블 공간을 가진 모델이 하나의 기초 모델로 일반화될 수 있는가?
• RQ2 대규모 텍스트 속성 그래프에서 자기지도 사전 훈련에 대해 마스킹 그래프 모델링(MGM)이 얼마나 효과적인가?
• RQ3 그래프 지시 훈련이 새로운 그래프 작업과 클래스에 대한 영향 없는 예측을 얼마나 잘 가능하게 하는가?
• RQ4 소수의 샘플에서의 전이 설정과 영향 없는 전이 설정에서 계층적 LM-GNN 아키텍처는 표준 GNN과 비교하여 어떻게 성능을 내는가?
• RQ5 사전 훈련 데이터의 다양성과 모델 구성 요소 선택(예: LM, GNN)이 최종 성능에 미치는 영향은 어떠한가?

주요 결과

• UniGraph는 5개 영역에서 온 11개 데이터셋의 노드, 간선, 그래프 수준 작업 전반에서 전반적인 베이스라인(특수 작업 GNN 및 Prodigy, OFA와 같은 다중 작업 사전 훈련 모델 포함)을 모두 압도한다.
• 1,110만 개 노드를 포함한 ogbn-Papers100M 데이터셋에서 UniGraph는 어떤 미세조정 없이도 자기지도 표현 학습에서 최신 기술 수준의 성능을 달성한다.
• 영향 없는 전이 설정에서 UniGraph는 Llama2 및 Vicuna와 같은 오픈소스 LLM보다 뚜렷이 뛰어난 성능을 보이며, 그래프 임베딩과 자연어 간의 효과적인 정렬을 입증한다.
• 절단 실험 결과 모든 구성 요소—MGM, 잠재 손실, PPR 샘플링, GNN—이 성능 향상에 기여하며, 특히 PPR 샘플링이 다양한 영역 간의 일반화를 향상시킨다.
• 더 큰 LM(예: DeBERTa-V3-base)과 더 깊은 GNN은 성능 향상을 이끌지만 추론 속도에 영향을 주므로 적절한 하드웨어를 사용할 경우 확장성이 가능함을 시사한다.
• 학습 시간과 메모리 비용은 MLM 기반 표준 LLM 사전 훈련과 유사하여, LM 구성 요소가 계산 비용을 주로 차지함을 보여주며, 이는 대규모 그래프에 대해 효율적인 프레임워크임을 의미한다.