KGPFN: Unlocking the Potential of Knowledge Graph Foundation Model via In-Context Learning

KGPFN 기술 요약: 인-컨텍스트 학습을 통한 지식 그래프 기반 모델의 잠재력Unlocking

1. 문제 정의

본 논문은 현재 지식 그래프 (KG) 기반 모델의 중요한 한계를 다룹니다. 즉, 이러한 모델들은 전이 가능한 관계 패턴 (보편성) 학습에는 탁월하지만, 기반 모델의 두 번째 기둥인 인-컨텍스트 학습 (ICL) 은 largely 소홀히 한다는 점입니다. 기존 KG 기반 모델들은 광범위하고 재사용 가능한 패턴을 과도하게 적용하는 경향이 있으며, 미세 조정 (fine-tuning) 없이 특정 그래프 인스턴스에 대한 예측을 적응시키지 못합니다.

저자들은 효과적인 KG 추론이 본질적으로 구조화되고 이질적인 두 가지 유형의 컨텍스트에 조건부적으로 의존해야 한다고 주장합니다.

• 로컬 컨텍스트: 추론이 발생하는 쿼리 특정 이웃입니다. 동일한 구조적 단서 (예: "삼각형" 패턴) 는 주변 서브그래프 (예: 출생지 대 국적) 에 따라 다른 결론을 낳을 수 있습니다.
• 글로벌 컨텍스트: 여러 인스턴스에 걸쳐 축적된 관계 수준의 증거로, 특정 관계가 그래프 전반에서 일반적으로 어떻게 행동하는지 설명합니다 (예: 공존하는 관계, 경로 빈도, 유효한 엔티티에 대한 제약 조건).

현재의 방법론들은 이러한 상호 보완적인 신호들을 통합하여 재학습 없이 보이지 않는 그래프에 대한 제로샷 (zero-shot) 또는 퓨샷 (few-shot) 적응을 가능하게 하는 데 어려움을 겪고 있습니다.

2. 방법론: KGPFN

저자들은 구조화된 KG 데이터에서 인-컨텍스트 학습을 수행하기 위해 Prior-Data Fitted Network (PFN) 패러다임을 활용하는 KG 기반 모델인 KGPFN을 제안합니다. 이 프레임워크는 대상 그래프에 대한 미세 조정이 없이 사전 학습과 추론 시 적응에 의존하여 작동합니다.

주요 구성 요소:

1. 컨텍스트 구성:

   • 글로벌 컨텍스트: 쿼리 관계 $r$에 대해, 모델은 그래프에서 양의 삼중항 (positive triples) 과 음의 삼중항 (negative triples) 인스턴스 집합을 검색합니다. 이러한 인스턴스와 그들의 로컬 이웃은 관계 수준의 규칙성을 포착하는 서포트 집합을 형성합니다.
   • 로컬 컨텍스트: 모든 후보에 대해 전체 서브그래프를 추출하는 것 (계산 비용이 많이 듦) 대신, 모델은 쿼리 헤드 엔티티 주변에 헤드 중심 $k$-홉 이웃을 구성합니다. 이는 대규모 후보 집합으로 효율적으로 확장됩니다.

2. 컨텍스트 그래프 인코더:

   • 관계 표현: 관계 그래프(노드는 관계 유형이고 에지는 상호작용을 나타냄) 에서의 메시지 전달 메커니즘을 통해 쿼리 조건부 관계 임베딩을 생성합니다. 이는 그래프 간 관계 불변성을 포착합니다.
   • 로컬 인코딩: 다층 NBFNet(Neural Belief Propagation Network) 이 헤드 중심 서브그래프를 인코딩하여 최대 $k$홉까지의 증거를 집계하는 멀티스케일 로컬 컨텍스트 벡터를 생성합니다.
   • 테일 인식 점수 매기기: 모델은 상호작용 특징 (TransE, DistMult, 코사인 유사도) 을 사용하여 헤드 표현, 관계 임베딩, 테일 표현 (마지막 NBFNet 레이어에서) 을 결합하여 테일 인식 점수 표현을 형성합니다.
   • 퓨전: 이질적인 특징 (로컬 컨텍스트, 관계 특징, 상호작용 특징) 은 경량 어댑터를 통해 공유 잠재 공간에 정렬된 후 연결됩니다.

3. PFN 아키텍처 (인-컨텍스트 학습):

   • 쿼리와 레이블이 지정된 글로벌 컨텍스트 집합의 퓨전된 표현이 PFN 모듈로 입력됩니다.
   • 특징 어텐션: 각 샘플 (로컬 컨텍스트) 내의 특징 차원을 정제합니다.
   • 샘플 어텐션: 컨텍스트 예제들 (글로벌 컨텍스트) 간 증거를 집계합니다. 이를 통해 쿼리는 레이블이 지정된 서포트 집합에서 관계별 증거를 검색할 수 있습니다.
   • 메커니즘: PFN은 베이지안 사후 예측을 근사화합니다. 컨텍스트 집합을 관측된 증거로 간주하여 쿼리에 대한 예측 분포를 유도함으로써, 언제 재사용 가능한 패턴을 인스턴스화하고 언제 컨텍스트 증거에 따라 이를 무시할지 학습합니다.

4. 학습 목적:

   • 모델은 양과 음을 분리하기 위한 이진 교차 엔트로피 (BCE) 손실과 샘플링된 음보다 양의 순위를 강제하기 위한 소프트맥스 교차 엔트로피 손실의 조합을 사용하여 다양한 KG 에서 학습됩니다.
   • 적응형 온도 메커니즘을 통해 가중치가 부여된 하드 네거티브를 사용한 네거티브 샘플링이 적용됩니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 제시합니다.

• KG 를 위한 듀얼 컨텍스트 구성: 저자들은 전이 가능한 패턴에 주로 초점을 맞춘 기존 기반 모델의 간극을 해결하기 위해 로컬 (쿼리 특정 이웃) 과 글로벌 (관계 수준 규칙성) 컨텍스트를 명시적으로 구축하는 전략을 도입합니다.
• PFN 기반 인-컨텍스트 학습: KGPFN 은 구조화된 KG 데이터에 Prior-Data Fitted Network 패러다임을 적용한 최초의 모델입니다. 특징 수준 및 샘플 수준 어텐션을 결합함으로써 파라미터 업데이트 없이 보이지 않는 그래프에 대한 추론 시 적응을 가능하게 합니다.
• 최첨단 성능: 광범위한 실험을 통해 KGPFN 이 유도 (inductive) 및 완전 유도 (full-inductive) 설정 모두에서 우수한 성능을 달성하며, 인-컨텍스트 학습만 사용하는 경쟁적인 미세 조정 베이스라인을 능가함을 입증했습니다.

4. 실험 결과

모델은 전이 (Transductive), 유도 (Inductive, 보이지 않는 엔티티), 완전 유도 (Fully Inductive, 보이지 않는 엔티티 및 관계) 의 세 가지 설정에서 57 개 KG 벤치마크에 대해 평가되었습니다.

• 성능: KGPFN 은 모든 57 개 데이터셋에서 평균 평균 역순위 (MRR) 와 Hits@10 에서 최고의 성과를 거두었습니다.
• 유도 능력: 가장 까다로운 완전 유도 설정에서 KGPFN 의 우위가 가장 두드러졌으며, ULTRA, TRIX, MOTIF 와 같은 모델들이 대상 데이터셋에 대해 미세 조정되었음에도 불구하고 일관되게 능가했습니다.
• 컨텍스트 민감도: 분석 결과, 글로벌 컨텍스트 내 음의 인스턴스 수를 늘리는 것이 성능을 크게 향상시키는 반면, 양의 인스턴스 수는 영향이 적음이 확인되었습니다. 또한 멀티홉 (3 홉) 로컬 컨텍스트를 사용하는 것이 결정적이었으며, 1 홉 또는 2 홉 이웃은 성능 저하로 이어졌습니다.
• 시각화: 어텐션 맵은 모델이 올바른 테일 후보에 대해서는 올바른 컨텍스트 인스턴스에 어텐션을 집중시키고, 잘못된 후보에 대해서는 음의 인스턴스 위로 어텐션을 분산시키는 것을 확인시켜 주었습니다.

5. 중요성과 주장

본 논문은 KGPFN 이 지식 그래프 추론을 위한 컨텍스트 기반 적응의 힘을 입증한다고 주장합니다. PFN 패러다임을 통해 글로벌 관계 규칙성과 로컬 구조적 단서를 통합함으로써, 모델은 이전에 보이지 않았던 그래프에 대한 강력한 일반화를 달성합니다.

저자들은 그들의 접근 방식이 모델이 구조적 불변성뿐만 아니라 기능적 표현력 (functional expressivity) 을 학습할 수 있게 한다고 강조합니다. 이론적으로 PFN 어텐션 메커니즘은 인-컨텍스트 서포트 집합에서 추론된 레이블 가중치 패턴 통계와 로컬 구조적 증거를 정렬함으로써 관계별 점수 매기기 함수를 즉석에서 구성합니다. 이는 미세 조정 없이 보이지 않는 관계에 적응할 수 있게 하여, 진정한 범용 KG 기반 모델을 향한 중요한 한 걸음을 의미합니다.

이 연구는 사전 출판 (preprint) 으로 발표되었으며, 코드는 커뮤니티에 공개되었습니다. 저자들은 컴퓨팅 자원 제약으로 인해 확장 법칙 (scaling laws) 검증과 관련된 한계를 지적합니다.