How to Make LMs Strong Node Classifiers?

기술 요약: AUGLM – 노드 분류를 위한 증강 그래프 언어 모델(Augmented Graph Language Model)

문제 정의
언어 모델(LM)은 그래프 학습 작업에서 유망한 가능성을 보여주었으나, 텍스트 속성 그래프(Text-Attributed Graphs, TAGs)의 노드 분류를 위한 기존 접근 방식들은 성능과 아키텍처 유연성 사이에서 근본적인 트레이드오프 문제에 직면해 있습니다. InstructGLM 및 LLaGA와 같은 최첨단(SOTA) 방법들은 각 노드에 고유한 토큰을 할당하기 위해 어휘 사전(vocabulary)을 확장하거나 토큰 공간에 위상 인식 임베딩을 주입하는 등 LM의 아키텍처를 수정함으로써 강력한 결과를 달로 성취합니다. 이러한 방식은 효과적이긴 하지만, 백본 LM의 범용성을 훼손하여 다양한 데이터셋에 걸친 원활한 멀티태스크 학습을 방해하며, 아키텍처 접근 권한이 없는 독점적 또는 폐쇄형 모델의 사용을 제한합니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: 어떻게 하면 별도의 아키텍처 수정 없이도 기성(off-the-shelf) 텍스트-투-텍스트 지시어 튜닝된 LM이 경쟁력 있는 노드 분류 성능을 달성할 수 있는가?

방법론: AUGLM
저자들은 아키텍처 변경 대신 데이터 증강 전략을 통해 표준 LM을 강화하는 프레임워크인 AUGLM(Augmented Graph Language Model)을 제안합니다. AUGLM은 LM의 원래 구조를 보존하면서 입력 컨텍스트를 풍부하게 함으로써 노드 분류를 텍스트-투-텍스트 작업으로 재정의합니다. 이 프레임워크는 세 가지 핵심 구성 요소에 의존합니다:

1. 검색 기반 집계(Retrieval-Based Aggregation):
   단순한 멀티 홉 이웃 집계(노이즈를 유발하거나 "이웃 폭발"을 일으킬 수 있음)에 의존하는 대신, AUGLM은 풍부한 컨텍스트 정보를 제공하기 위해 두 가지 상호 보완적인 검색 전략을 채택합니다:

   • 위상적 검색(Topological Retrieval): 고정된 홉(hop) 거리가 아닌, 그래프 구조를 기반으로 타겟 노드에 가장 관련성이 높은 이웃을 식별하기 위해 개인화 페이지랭크(Personalized PageRank, PPR)를 활용합니다. 이러한 상위 K개의 PPR 이웃의 텍스트 특징을 결합하여 구조적 컨텍스트를 형성합니다.
   • 원형적 의미론적 검색(Prototypical Semantic Retrieval): 검색 증강 생성(RAG) 및 그래프 트랜스포머에서 영감을 받은 이 모듈은 타겟 노드가 아닌 모든 노드의 텍스트 특징을 외부 코퍼스로 취급합니다. 먼저 경량 GNN을 학습시켜 노드 예측을 생성함으로써 "프로토타입(prototypes)"—각 클래스를 대표하는 예시들—을 추출합니다. 각 클래스에 대해 신뢰도가 가장 높은 예시들이 프로토타입으로 선택됩니다. 시스템은 타겟 노드의 텍스트와 가장 의미론적으로 유사한 프로토타입(및 그와 관련된 PPR 이웃)을 검색합니다.

2. 분류기 가이드(Classifier Guidance, 후보군 프루닝):
   그래프 위상에 대한 LM의 내재된 이해력 한계를 해결하기 위해, AUGLM은 전문 모델의 가이드를 통합합니다. 사전 학습된 경량 GNN은 각 노드에 대해 상위 $I$개의 후보 레이블을 예측합니다. 이 후보들은 텍스트로 변환되어 프롬프트에 포함됩니다. 이 전략은 탐색 공간을 효과적으로 프루닝(pruning)하여, LM이 전체 레이블 공간이 아닌 위상 정보가 반영된 고도로 관련성 있는 후보들 사이를 구별하는 추론에 집중할 수 있도록 합니다.

3. 학습 및 최적화:
   프레임워크는 프로토타입과 후보를 생성하기 위한 경량 GNN($\psi$), 의미론적 리트리버($\phi$), 그리고 백본 LM($\theta$)이라는 세 가지 매개변수화된 모듈을 포함합니다:

   • GNN은 독립적으로 사전 학습됩니다.
   • 의미론적 리트리버는 검색 확률을 정답 레이블을 생성할 LM의 가능도(likelihood)와 일치시키기 위해 분포 매칭 손실(KL 발산)을 사용하여 학습됩니다.
   • 백본 LM은 증강된 입력 프롬프트에 대해 표준 토큰 단위 음의 로그 가능도(NLL) 손실을 사용하여 파인튜닝됩니다.
   • 결정적으로, 리트리버와 LM의 계산 그래프는 특정 학습 단계에서 stop-gradient 연산자를 통해 분리되며, 이를 통해 효율적인 역전파를 보장합니다.

주요 기여

• 아키텍처 보존: 기존의 SOTA 방법들과 달리, AUGLM은 LM의 아키텍처나 어휘 사전에 대한 수정이 필요하지 않으므로, 어떤 기성 텍스트-투-텍스트 LM도 사용할 수 있습니다.
• 이중 전략 증강: 단순한 이웃 집계의 한계를 극복하기 위해 위상적(PPR) 정보와 의미론적(프로토타입 기반) 정보를 결합한 하이브리드 검색 메커니즘을 도입했습니다.
• 경량 가이드: GNN으로부터 전달된 위상적 귀납 편향(inductive bias)을 후보 레이블 프루닝을 통해 LM으로 전달하는 새로운 방법을 제시하여, LM의 내부 파라미터를 변경하지 않고도 성능을 향상시켰습니다.
• 멀티태스크 유연성: 이 프레임워크는 데이터셋별 토큰 확장이 필요 없이 다양한 데이터셋에 대한 공동 학습을 지원하며, 이는 아키텍처를 수정하는 접근 방식에서 흔히 상실되는 능력입니다.

실험 결과
저자들은 Cora, Pubmed, ogbn-arxiv, ogbn-products의 네 가지 실제 TAG 데이터셋에서 AUGLM을 평가했습니다.

• 성능: Flan-T5 백본을 탑재한 AUGLM은 일관되적으로 SOTA 텍스트 출력 노드 분류기(InstructGLM 및 LLaGA 포함)보다 우수한 성능을 보였습니다.
• 벡터 출력 모델과의 비교: AUGLM은 순수하게 텍스트 도메인에서 작동함에도 불구하고, 최고의 벡터 출력 모델(TAPE 및 GLEM+RevGAT 등)과 대등하거나 일부 사례(예: ogbn-products)에서는 이를 능가하는 성능을 달성했습니다.
• 절제 연구(Ablation Studies): 실험을 통해 위상적 검색, 의미론적 검색, 레이블 프루닝의 모든 구성 요소가 성능에 긍정적으로 기여함을 확인했으나, 그 상대적 중요도는 데이터셋마다 달랐습니다.
• 공동 학습: 네 가지 데이터셋 모두에서 공동 학습된 모델은 각 데이터셋에 대해 개별적으로 학습된 모델과 대등한 성능을 달성하였으며, 이는 다양한 그래프 구조에 걸쳐 성능 저하 없이 일반화할 수 있는 프레임워크의 능력을 입증합니다.

의의 및 주장
본 논문은 AUGLM을 특화된 태스크 중심의 노드 분류기와 범용 파운데이션 모델 사이의 간극을 메우는 중요한 단계로 포지셔닝합니다. 고성능 그래프 학습이 아키텍처 수정이 아닌 입력 증강을 통해 달성될 수 있음을 입증함으로써, 본 연구는 더욱 다재다능하고 널리 적용 가능한 그래프 학습 모델을 향한 경로를 제시합니다. 저자들은 이 접근 방식이 지시어 튜닝(instruction tuning)의 이점(특히 다양한 데이터셋에 대한 공동 학습 능력)을 유지하면서도, 기존의 경직된 그래프-LM 통합 방식의 한계를 극복한다고 주장합니다. 결론적으로, 현재의 LM이 그래프 구조를 이해하는 데 있어 내재적인 한계가 있지만, 전략적인 데이터 증강을 통해 이를 효과적으로 완화할 수 있으며, 이를 통해 여러 그래프 학습 작업에서 탁월한 성능을 발휘하는 통합 모델의 길을 열 수 있다고 밝히고 있습니다.