기술 요약: 임베딩 기반 링크 예측의 이론적 한계에 대하여

1. 문제 정의

본 논문은 현대 지식 그래프 임베딩(KGE) 모델의 근본적인 한계인 선형 출력층에 의한 **랭크 병목 현상(rank bottleneck)**을 다룹니다. 표준 KGE 아키텍처에서 엔티티와 관계는 저차원 임베딩(통상적으로 $d \in [10^2, 10^4]$)으로 매핑된 후, 선형층을 통해 고차원 엔티티 공간(통상적으로 $|E| \in [10^5, 10^9]$)으로 투영됩니다.

저자들은 이러한 선형 투영이 모델의 출력 점수를 $|E|$ 차원 공간 내의 $d$ 차원 선형 부분 공간으로 국한시킨다고 주장합니다. $d \ll |E|$일 때, 이는 모델의 표현력을 심각하게 제한하는 "랭크 병목"을 생성합니다. 구체적으로, 모델은 상위 신경 인코더(예: 그래프 신경망 또는 트랜스포머)의 성능과 무관하게, 복잡한 다운스트림 추론 작업에 필요한 임의의 순위 순서, 부호 패턴(양수/음수 트리플), 또는 확률 분포를 표현할 수 없습니다.

2. 방법론

이론적 분석

저자들은 선형 디코더 KGE(LD-KGE)가 세 가지 별개의 예측 작업을 충족하기 위해 필요한 임베딩 차원 $d$의 필요 경계치를 도출합니다:

1. 순위 재구성 (Ranking Reconstruction, RR): 쿼리에 대한 객체들을 올바르게 정렬함.
2. 부호 재구성 (Sign Reconstruction, SR): 트리플을 참 또는 거짓으로 올바르게 분류함.
3. 분포 재구성 (Distributional Reconstruction, DR): 트리플에 대해 보정된 확률을 할당함.

이들은 두 가지 조건 하에서 이 작업들을 분석합니다:

• 보편성 (Universality): 주어진 크기의 그래프에 대해 가능한 모든 패턴을 표현할 수 있는 능력. 논문은 보편성을 위해 $d$가 적어도 $|E|-1$ (RR 및 DR의 경우) 또는 $|E|$ (SR의 경우)여야 함을 증명합니다. 이러한 경계치는 대규모 그래프에서는 실질적으로 불가능합니다.
• 실현 가능성 (Realisability): 관찰된 특정 그래프와 일치하는 패턴을 표현할 수 있는 능력. 저자들은 필요 차원이 그래프의 인접 행렬(또는 부호 행렬)의 **부호 랭크(sign rank)**에 달려 있음을 보여줍니다. 부호 랭크는 일반적으로 전체 랭크보다 작지만, 이를 계산하는 것은 NP-hard입니다.

실용적인 충분 조건을 제공하기 위해, 저자들은 부호 랭크를 그래프 연결성과 연결합니다. 그들은 필요 차원이 그래프의 연결성(구체적으로 양수 항목의 연속적인 블록 수)에 의해 제한됨을 입증하며, 엔티티를 재정렬하기 위해 역 커틸-맥키(Reverse Cuthill-McKee, RCM) 휴리스틱을 사용합니다. 이를 통해 필요한 차원이 단순히 전체 엔티티 수가 아니라 그래프의 연결성에 따라 스케일링됨을 보여줍니다. 그러나 최적의 정렬을 사용하더라도, 밀집된 그래프는 표준적인 관행을 훨씬 초과하는 차원을 요구합니다.

제안된 솔루션: KGE-MOS

차원 $d$를 높이는 막대한 파라미터 비용을 들이지 않고 이 병목 현상을 해결하기 위해, 저자들은 KGE-MOS(Mixture of Softmaxes를 이용한 지식 그래프 임베딩)를 제안합니다.

• 아키텍처: 단일 선형층과 소프트맥스를 사용하는 대신, KGE-MOS는 출력층을 $K$개의 소프트맥스 혼합(mixture)으로 대체합니다.
• 메커리즘: 쿼리 $(s, r)$에 대해, 모델은 $K$개 분포의 혼합을 계산합니다:
  $$P(O|s, r) = \sum_{k=1}^K \pi_k(h_{s,r}) \cdot \text{softmax}(f_k(h_{s,r})E^\top)$$
  여기서 $h_{s,r}$은 은닉 상태, $E$는 공유 엔티티 임베딩 행렬, $f_k$는 구성 요소별 비선형 투영(MLP), 그리고 $\pi_k$는 혼합 가중치입니다.
• 효율성: 엔티티 임베딩 $E$는 모든 구성 요소 간에 공유됩니다. 추가되는 파라미터는 투영 가중치 $f_k$와 혼합 가중치 $\pi_k$로 제한되며, 결과적으로 $O(Kd^2)$의 비용이 발생하여 $d$를 높이는 데 드는 $O(d|E|)$ 비용에 비해 매우 미미합니다.

3. 주요 기여

1. 이론적 경계치: 본 논문은 보편성 및 실현 가능성과 관련하여 모든 LD-KGE(이차 형식 및 신경망 모두)에 대한 최초의 필요 경계치를 제공합니다. 이는 선형 출력층이 대규모, 밀집 그래프에서 표현력을 근본적으로 제한함을 확립합니다.
2. 연결성 분석: 그래프 연결성(부호 변화)을 기반으로 한 실현 가능성을 위한 충분 차원을 도출하여, 밀집된 그래프가 희소한 그래프보다 이론적으로 더 맞추기 어렵다는 것을 보여줍니다.
3. KGE-MOS: 랭크 병목을 효율적으로 깨뜨리는 드롭인(drop-in) 비선형 출력층을 도입합니다.
4. 실험적 검증: KGE-MOS가 대규모 밀집 데이터셋에서 기존의 LD-KGE들보다 순위 정확도(MRR)와 확률적 적합도(NLL) 모두에서 개선된 성능을 보임을 입증합니다. 또한, 단순히 임베딩 차원을 높이는 것이 훈련 불안정성(예: 배치 크기 제약)으로 인해 수확 체감의 법칙을 따름을 보여줍니다.

4. 실험 결과

저자들은 다섯 가지 데이터셋(FB15k-237, Hetionet, ogbl-biokg, openbiolink, WN18RR)에 대해 네 가지 베이스라인 모델(DISTMULT, COMPLEX, RESCAL, CONVE) 및 하나의 신경 인코더(COMPGCN)를 사용하여 KGE-MOS를 평가했습니다.

• 대규모/밀집 그래프에서의 성능: 가장 크고 밀집도가 높은 세 가지 데이터셋(Hetionet, ogbl-biokg, openbiolink)에서 KGE-MOS는 MRR과 NLL 모두에서 베이스라인 LD-KGE들을 지속적으로 능가했습니다. 성능 향상은 이론적 병목이 가장 심각한 고연결성 영역에서 가장 두드러졌습니다.
• 소규모 그래프에서의 성능: 더 작은 FB15k-237 데이터셋에서 KGE-MOS는 무시할 만한 개선 또는 약간의 성능 저하를 보였으며, 이는 병목 현상이 규모가 작고 덜 밀집된 그래프에서는 덜 결정적임을 확인시켜 줍니다.
• 차원 스케일링: 임베딩 차원 $d$를 200에서 1000으로 높이면 베이스라인의 성능은 향상되었으나, GPU 메모리 제약으로 인해 배치 크기를 줄여야 했고, 이는 궁극적으로 성능을 해쳤습니다. 반면, $d=200$에서의 KGE-MOS는 훨씬 적은 파라미터로 $d=1000$인 베이스라인과 대등한 성능을 달성했습니다.
• 절제 연구(Ablation): 소프트맥스 구성 요소($K$)를 늘릴수록 성능이 일관되게 향상되었습니다. 비선형 투영 $f_k$ 또한 이득에 기여했지만, 병목을 깨뜨리는 주된 동력은 혼합(mixture) 자체였습니다.
• 랭크 타파: 로그 확률 행렬의 경험적 랭크 분석 결과, 표준 모델은 $d+1$ 랭크로 제한되었던 반면, KGE-MOS는 엔티티 집합의 전체 차원에 근접하는 랭크를 달축하여 효과적으로 병목을 깨뜨렸습니다.

5. 의의 및 주장

본 논문은 링크 예측 분야가 인코더를 개선하는 데(예: GNN 또는 트랜스포머 사용) 집중해 온 반면, 출력층의 한계는 간과해 왔다고 주장합니다. 저자들은 더 표현력이 풍부한 비선형 출력층을 탐구하는 것이 대규모 및 밀집 그래프로 KGE를 확장하기 위한 필수적이고 유망한 경로라고 주장합니다.

그들은 K-MOS가 낮은 파라미터 비용으로 성능과 확률적 충실도를 개선하는 "드롭인" 솔루션을 제공하며, 단순히 임베딩 차원을 높이는 것이 더 큰 그래프를 처리하는 유일한 방법이라는 가설에 도전한다고 강조합니다. 이 연구는 향-임베딩 모델의 디코딩 과정에 존재하는 구조적 한계를 해결하기 위해, 미래의 연구가 인코더 중심의 개선을 넘어 출력층으로 나아가야 함을 시사합니다.