GRATE: Temporal Extensions for Inductive KG Foundation Models via Gated Rotary Attention

기술 요약: GRATE: 게이트형 로터리 어텐션을 통한 유도적 KG 파운데이션 모델의 시간적 확장

1. 문제 정의

본 논문은 유도적 지식 그래프(Inductive Knowledge Graph, KG) 파운데이션 모델(ULTRA 및 TRIX와 같은)을 **시계열 지식 그래프(Temporal Knowledge Graph, TKG)**로 확장하는 과제를 다룹ون다. 기존의 파운데이션 모델들은 어휘 집합에 의존하지 않는 표현을 학습하여 보지 못한 엔티티와 관계에 일반화함으로써 강력한 유도적 전이(inductive transfer)를 달성하지만, 시간적 역동성을 처리할 메커니즘이 부족하다.

현재의 TKG 모델들은 일반적으로 엔티티, 관계, 타임스탬프를 위한 데이터셋 특정 임베딩을 학습하는 트랜스덕티브(transductive) 설정에 의존한다. 이는 새로운 어휘(엔티티, 관계 또는 시점)가 지속적으로 유입되는 시나리오(예: 급변하는 정치적 행위자 또는 빠르게 진화하는 이벤트 스트림)에는 부적합하다. 또한, 기존의 유도적 KG 모델들은 모든 타임스탬프의 사실들을 쿼리와 동일하게 정렬된 것으로 취급하며, 지원하는 사실의 관련성이 쿼리 시간으로부터의 **상대적 시간 변위(relative temporal displacement)**에 따라 달라질 수 있다는 점을 고려하지 못한다. 저자들은 파운데온 모델의 파라미터 프리(parameter-free) 전이 가능성을 희생하지 않으면서, 유도적 시간 추론(서로 다른 어휘와 타임스탬프에 대한 일반화)을 가능하게 하는 방법론이 없다는 격차를 식별하였다.

2. 방법론: GRATE

저자들은 NBFNet 스타일의 KG 파운데이션 모델을 증강하기 위해 설계된 파라미터 프리 메시 함수인 GRATE(Gated Rotary Attention for Temporal Encoding)를 제안한다. GRATE는 엔티티 측 메시 전달 메커니즘에 통합되되, 학습 가능한 파라미터를 도입하지 않음으로써 서로 다른 어휘를 가진 그래프 간의 전이 능력을 보존한다.

GRATE는 두 가지 상호 보완적인 메커니즘을 통해 작동한다:

A. 로터리 상대적 시간 인코딩 (Rotary Relative-Time Encoding)
학습된 타임스탬프 임베딩을 사용하는 대신, GRATE는 상대적 시간 차이를 통해 시간을 인코딩한다. 쿼리가 시간 $\tau$에 있고 관측된 에지 $(u, r, v)$가 시간 $t_{uv}$에 존재할 때, 상대적 시간 간격은 $\Delta t_{uv} = \tau - t_{uv}$로 정의된다.

• 이 방법은 Transformer에서 유래한 **로터리 위치 임베딩(Rotary Position Embeddings, RoPE)**을 변형하여, $\Delta t_{uv}$에 따라 구조적 메시 $m_{uv}$를 회전시킨다.
• 회전 행렬 $R(\Delta t_{uv})$는 메시 벡터에 주파수 종속적 회전을 적용한다. 이는 구조적 메시를 시간적 변위에 대한 연속 함수로 변환하여, 모델이 절대적인 타임스탬프 정체성과 무관하게 쿼리와 시간적으로 가까운 사실과 먼 사실을 구분할 수 있게 한다.

B. 쿼리 조건부 게이팅 (Query-Conditioned Gating)
회전만으로는 어떤 시간적으로 이동된 사실이 특정 쿼리에 관련이 있는지를 결정할 수 없다. GRATE는 쿼리와의 정렬(alignment)을 기준으로 메시의 가중치를 결정하는 게이팅 메커니즘을 도입한다.

• 쿼리의 헤드 엔티티 표현과 쿼리 관계 임베딩의 합을 통해 쿼리 상태 $Q_q$를 구성한다.
• 스칼라 게이트 $g_{uv}$는 시그모이드 함수를 사용하여 계산된다: $g_{uv} = \sigma(Q_q^\top m_{uv}^{rot} / \sqrt{d})$, 여기서 $m_{uv}^{rot}$는 회전된 메시이다.
• 이 게이트는 각 에지를 독립적으로 평가하여, 쿼리에 대해 시간적 및 의미적으로 모두 관련 있는 신호를 강조하고 관련 없는 신호는 억제한다. Softmax 어텐션과 달리, 이러한 독립적 게이팅은 이웃의 크기가 크게 변할 수 있는 유도적 전이 상황에서도 안정적이다.

C. 멀티-TKG 사전 학습 (Multi-TKG Pretraining)
여러 TKG를 동시에 학습시키기 위해, 저자들은 각 소스 데이터셋의 타임스탬프를 조인 학습 전에 서로 겹치지 않는 정수 범위로 재매핑한다. GRATE는 오직 상대적 간격($\Delta t$)에만 의존하므로, 이 재매핑은 데이터셋 간의 ID 충돌을 방지하면서도 각 소스별 시간적 관계를 보존한다.

3. 주요 기여

1. GRATE 메커니즘: 상대적 시간 회전과 쿼리 조건부 게이팅을 통해 유도적 KG 파운데이션 모델이 시간적 추론을 수행할 수 있도록 하는 새로운 파라미터 프리 시간 메시 함수를 제안한다.
2. 유도적 전이 벤치마크: 훈련 그래프와 추론 그래프 간의 엔티티, 관계, 타임스탬프가 서로 겹치지 않는 보간(interpolation) 및 외삽(extrapolation) 체제를 모두 포함하는 벤치마크 스위트인 GDELTINDT와 WIKIINDT를 구축하였다.
3. 제로샷 시간 추론: 단일 공동 사전 학습된 GRATE 체크포인트가 미세 조정 없이도 보지 못한 데이터셋에 대해 시간적 추론 성능을 향상시킬 수 있음을 입증하였으며, 정적 베이스 모델 및 기존 유도적 베이스라인보다 우수한 성능을 보였다.

4. 실험 결과

저자들은 단일 공동 사전 학습된 체크포인트를 사용하여 세 가지 체제에서 GRATE를 평가하였다:

• 트랜스덕티브 인-디스트리뷰션 (ICEWS14, ICEWS05-15): ULTRA와 TRIX에 GRATE를 추가하면 성능이 크게 향상된다. MRR에서 +25% ~ +38%, Hits@1에서 **+38% ~ +58%**의 이득을 보였으며, 이는 상대적 시간 정렬이 주로 상위 순위 정밀도를 개선함을 나타낸다.
• 유도적 전이 (GDELTINDT, WIKIINDT):
  • GDELT(밀집 이벤트 그래프)에서 TRIX + GRATE는 보간에서 +0.16 ~ +0.21 MRR, 외삽에서 +0.12 ~ +0.18 MRR만큼 재귀 전용(recurrence-only) 베이스라인을 능가한다. 게이팅 메커니즘은 풍부한 이웃 사실들 사이에서 효과적으로 변별력을 발휘한다.
  • WIKI(고립된 단언들의 희소 그래프)에서는 이득이 더 가변적이다. "리터럴 재귀(literal recurrence)"(정확히 1년 전의 사실이 반복됨)가 높은 경우, GRATE는 이러한 일치 항목에 고정되어 상당한 이득을 얻는다. 그러나 재귀가 낮은 설정에서는 정적 베이스 모델이 GRATE보다 나은 경우가 있으나, GRATE는 여전히 일반적으로 재귀 베이스라인보다는 우수하다.
• 제로샷 전이 (ICEWS18, YAGO): GRATE는 ICEWS18의 Hits@10에서 LLM 기반 예측 방법(ICL, GenTKG)보다 우수한 성능을 보였으며, 텍스트 사전 지식 없이 순수하게 구조 중심임에도 불구하고 YAGO에서 더 높은 Hits@1 및 MRR을 달성했다.
• 절제 연구 (Ablation Studies): 회전 또는 쿼리 조건부 게이트를 제거하면 성능이 저하된다. 절제 연구는 상대적 시간 정렬이 대부분의 시간적 신호를 포착하는 반면, 게이트는 특히 구조적 단서가 약한 제로샷 설정에서 추가적인 이득을 제공함을 확인해 준다.

5. 의의 및 주장

본 논문은 파라미터 프리 전이 가능성을 희생하지 않고도 유도적 KG 파운데이션 모델에 시간적 추론을 추가할 수 있다고 주장한다. GRATE는 모델이 데이터셋 특정 임베딩 대신 상대적 시간 차이와 쿼리 조건부 관련성에 의존함으로써, 완전히 보지 못한 엔티티, 관계, 타임스탬프를 가진 TKG에 일반화될 수 있음을 보여준다.

저자들은 GRATE의 영향력이 베이스 모델의 관계 그래프의 풍부함(예: ULTRA의 단일 패스 집계보다 TRIX의 반복적 정교화가 더 많은 이득을 얻음)에 따라 달라진다고 언급한다. 또한 한계점도 인정하는데, 현재 GRATE는 시간을 순수하게 상대적 인덱스 차이를 통해 인코딩하며 시간적 입도(granularity)나 불규칙한 간격을 명시적으로 모델링하지 않으므로, 입도가 거친 데이터셋에서는 성능이 약화될 수 있다. 향후 연구로는 혼합 입도 사전 학습 및 더 확장 가능한 전파 전략을 탐구하는 것이 제안되었다.