[논문리뷰]Fact Embedding through Diffusion Model for Knowledge Graph Completion

Xiao Long, Liansheng Zhuang, Aodi Li, Houqiang Li, and Shafei Wang. 2024. Fact embedding through 197 diffusion model for knowledge graph completion. In Proceedings of the ACM on Web Conference 2024, WWW ’24, page 2020–2029, New York, NY, USA.Association for Computing Machinery

Problem Statement

1. 복잡한 연결 패턴의 처리 한계
기존의 KG 임베딩(KGE) 모델들은 주로 엔티티와 관계를 연속 벡터 공간에 매핑하고, 스코어 함수를 정의하여 사실 간의 연결 패턴을 캡처하는 방식을 사용한다. 그러나 실제 KG에서의 연결 패턴은 매우 복잡하여, 명시적이고 효율적인 스코어 함수를 정의하기 어렵다. 이로 인해 기존 모델들은 모든 연결 패턴을 효과적으로 처리하지 못하고, 성능이 제한되는 한계가 있다.

2. 제한된 패턴 모델링 능력
일부 연구들은 여러 KGE 모델을 결합하거나 더 복잡한 스코어 함수를 설계하여 다양한 패턴을 모델링하려고 시도했으나, 여전히 특정한 패턴에만 집중하는 경향이 있다. 이러한 결합 모델들은 제한된 수의 패턴만을 포착할 수 있으며, KG의 모든 다양하고 복잡한 패턴을 완전히 모델링하기에는 부족하다.

3. 스코어 함수의 비효율성
기존 연구에서는 스코어 함수를 통해 사실의 타당성을 측정하는 방식이 주로 사용되었으나, 이러한 접근 방식은 KG의 복잡성을 반영하지 못하여 비효율적이다. 스코어 함수를 통해 모든 사실의 타당성을 효과적으로 측정하는 것은 사실상 불가능하며, 이는 KG 완성 작업에서 성능 저하로 이어진다.

4. 사실의 분포를 직접 학습하지 않음
기존 모델들은 연결 패턴을 명시적으로 정의하고 이를 기반으로 학습하는 방식을 택했으나, 타당한 사실의 분포를 직접적으로 학습하지 않는다. 이는 모델이 실제 KG에서의 사실들을 포괄적으로 이해하고 예측하는 데에 제한을 두게 된다.

Method

Architecture

Architecture of FDM

FDM은 두 개의 스테이지로 나눠진다. 첫 번째는 Forward diffusion process이고, 두 번째는 Reverse diffusion process이다.

1. forward diffusion process

• 사전에 정의된 노이즈 스케줄에 따라 알려진 사실에 점진적으로 가우시안 노이즈를 추가하여 시간 단계 \(T\)까지 진행.

2. reverse diffusion process

• Reverse 과정에서는 훈련된 조건부 사실 노이즈 제거기를 사용하여 가우시안 분포에서 벡터 공간 내 타당한 사실의 분포로 마르코프 전이(Markov Transition)를 모델링한다.
• 이 과정에서 명시적인 조건부 제약이 적용되며, 조건부 인코더가 조건 임베딩을 인코딩하고 연결 패턴을 학습하여 목표 사실 임베딩을 생성하고, 최종적으로 해당 엔터티 임베딩을 도출한다.

Forward Process

Forward 과정에서는 알려진 사실에 점진적으로 가우시안 노이즈를 추가하여, 점차적으로 사실이 노이즈에 의해 변형된다. 이 과정은 마르코프 체인을 통해 수행되며, 시간 단계 \(T\)까지 계속된다. 각 시간 단계에서, 사실의 임베딩에 노이즈를 추가하는 방식으로, 이는 사실을 점진적으로 더 복잡한 형태의 분포로 변환하는 역할을 한다.

1. 사실 임베딩(Fact Embedding) 정의
사실 \(\tau\)는 헤드 엔티티 \(h\), 릴레이션 \(r\), 그리고 테일 엔티티 \(t\)로 구성되며, 각각의 요소를 임베딩하여 하나의 벡터로 결합한 사실 임베딩 \(X_\tau\)를 얻는다. 이는 수식으로 다음과 같이 표현된다.

$$X_\tau = [X_h; X_r; X_t]$$

2. 노이즈 추가(Noise Addition)
각 시간 단계 \(t\)에서 사실 임베딩 \(X_\tau\)에 가우시안 노이즈를 추가하는 과정을 거친다. 이 과정은 사실을 점차적으로 노이즈화하여, 최종적으로는 완전히 노이즈화된 상태로 변환하는 것을 목표로 한다. 이를 수식으로 표현하면 다음과 같다.

$$q(X_{\tau_{t}} | X_{\tau_{t-1}}) = \mathcal{N}(X_{\tau_{t}}; \sqrt{1 - \beta_t} X_{\tau_{t-1}}, \beta_t I)$$

여기서 \(\beta_t\)는 시간 단계 \(t\)에서의 노이즈 추가 정도를 조절하는 변수이다. 이 과정은 \(t = 1\)부터 \(t = T\)까지 반복되며, 최종적으로 사실 임베딩은 완전히 노이즈화된 상태로 변하게 된다. 이 노이즈화된 임베딩은 Reverse 과정에서 원래의 사실로 복원된다. 이와 같이 Forward 과정은 사실을 점진적으로 노이즈화하여, 후속 과정에서 이를 복원하기 위한 기초를 마련하는 역할을 한다.

foraward process에서는 학습하는 파라미터가 없다. Forward 과정에서는 사실 임베딩에 점진적으로 가우시안 노이즈를 추가하는 과정만 이루어지며, 이 과정 자체는 고정된 노이즈 스케줄 \(\beta_t\)에 따라 수행된다. 즉, Forward 과정은 단순히 가우시안 노이즈를 추가하는 단계로, 이 과정에서 파라미터를 학습하지 않는다.

Backward Process

Reverse 과정에서는 노이즈화된 사실 임베딩을 원래의 타당한 사실 임베딩으로 복원하는 역할을 한다. 이 과정은 조건부 노이즈 제거기(Conditional Fact Denoiser)를 사용하여, 가우시안 분포에서 타당한 사실의 분포로 마르코프 전이를 모델링하는 방식으로 수행된다. 또한, 명시적인 조건부 제약을 적용하여, 사실 임베딩이 원래의 상태로 복원되도록 유도한다.

1. 조건부 인코더(Conditional Encoder) 정의
조건부 인코더는 알려진 조건 임베딩(즉, 헤드 엔티티 \(X_h\)와 릴레이션 \(X_r\))을 인코딩하여, 이들이 사실 복원 과정에서 가이드 역할을 할 수 있도록 한다. 조건부 인코더의 출력은 다음 수식으로 표현된다.

$$X_c = \text{ConditionalEncoder}(X_h, X_r)$$

2. 노이즈 제거 및 복원(Noise Removal and Restoration)
각 시간 단계 \(t\)에서 조건부 노이즈 제거기(Conditional Fact Denoiser)는 이전 단계의 노이즈화된 사실 임베딩 \(X_{\tau_t}\)을 받아, 이를 덜 노이즈화된 상태로 변환한다. 이 과정은 수식으로 다음과 같이 표현된다.

$$p_\theta(X_{\tau_{t-1}} | X_{\tau_t}, X_h, X_r) = \mathcal{N}(X_{\tau_{t-1}}; \mu_\theta(X_{\tau_t}, t, X_h, X_r), \sigma_t^2 I)$$

여기서 \(\mu_\theta\)는 조건부 노이즈 제거기에 의해 예측된 평균값이고, \(\sigma_t\)는 고정된 분산 값이다. \(\mu_\theta\)는 다음과 같이 다시 표현할 수 있다.

$$\mu_\theta(X_{\tau_t}, t, X_h, X_r) = \frac{1}{\sqrt{\alpha_t}} \left( X_{\tau_t} - \frac{\beta_t}{\sqrt{1 - \alpha_t}} \epsilon_\theta(X_{\tau_t}, t, X_h, X_r) \right)$$

여기서 \(\alpha_t\)는 시간 단계 \(t\)에서의 노이즈 스케줄을 의미하며, \(\epsilon_\theta\)는 조건부 노이즈 제거기에서 예측된 노이즈 값이다.

3. 조건부 제약 적용(Conditional Constraint Application)
Reverse 과정에서 생성된 사실 임베딩이 기존의 조건과 일치하도록 하기 위해 명시적인 조건부 제약이 적용된다. 이 제약은 다음과 같은 목적 함수로 정의되며, 사실 임베딩이 조건 임베딩과 최대한 일치하도록 유도한다.

$$F(X_{hr_t}, X_{hr}) = \| X_{hr_t} - X_{hr} \|^2_2$$

여기서 \(X_{hr_t}\)는 특정 시간 단계 \(t\)에서의 헤드 엔티티와 릴레이션의 임베딩을 의미하며, \(X_{hr}\)는 원래의 헤드 엔티티와 릴레이션의 임베딩이다. 최종적으로, 사실 임베딩은 다음과 같은 경사 하강법을 통해 업데이트된다:

$$X_{\tau_{t-1}} = X_{\tau_{t-1}} - \eta \nabla_{X_{hr_t}} F(X_{hr_t}, X_{hr})$$

여기서 \(\eta\)는 고정된 학습률을 나타낸다. 이 과정은 \(t = T\)부터 \(t = 1\)까지 역순으로 진행되며, 최종적으로 사실 임베딩은 원래의 타당한 상태로 복원된다. 이와 같이 Reverse 과정은 노이즈화된 데이터를 원래의 사실로 복원하여, KG 내에서 타당한 사실을 생성하는 역할을 한다.

Conditional Fact Denoiser

Conditional Fact Denoiser는 Reverse 과정에서 노이즈화된 사실 임베딩을 원래의 타당한 사실 임베딩으로 복원하는 핵심 역할을 한다. 이 모듈은 특히 KG의 특성을 고려하여 설계되었으며, MLP 기반의 간단하고 효율적인 구조를 채택하여 사실을 복원하는 과정을 수행한다.

1. 조건부 인코더(Conditional Encoder)
Conditional Encoder는 알려진 조건 임베딩(즉, 헤드 엔티티 \(X_h\)와 릴레이션 \(X_r\))을 인코딩하여, 이들이 사실 복원 과정에서 가이드 역할을 할 수 있도록 한다. 이는 수식으로 다음과 같이 표현된다.

$$X_c = \text{ConditionalEncoder}(X_h, X_r)$$

2. CFDenoiser 블록
CFDenoiser 블록은 입력된 노이즈화된 사실 임베딩 \(X_{\tau_t}\)와 시간 단계 임베딩 \(X_t\), 그리고 조건 임베딩 \(X_c\)을 받아, 이를 덜 노이즈화된 상태로 변환하는 과정에서 사용된다. 이는 다음과 같이 표현된다.

$$E = \text{CFDenoiserBlock}(X_{\tau_t}, X_t, X_c)$$

3. 노이즈 예측(Noise Prediction)
CFDenoiser 블록을 통해 계산된 중간 결과 \(E\)는 선형 계층을 통해 최종적으로 노이즈 예측값 \(\epsilon\)으로 변환된다. 이 과정은 다음 수식으로 표현된다.

$$\epsilon = \text{LinearLayer}(\text{LN}(E))$$

여기서 \(\text{LN}\)은 LayerNorm을 나타낸다. Conditional Fact Denoiser는 KG의 특성을 고려하여 간단하면서도 효율적인 MLP 기반의 구조를 채택하여, 노이즈화된 사실 임베딩을 원래의 상태로 복원하는 역할을 한다.

Training and Inference

Training과 Inference 단계에서는 Conditional Fact Denoiser를 사용하여 모델을 학습시키고, 실제로 데이터를 추론하는 과정을 수행한다. 이 과정에서 negative sampling 기법을 활용하여 모델을 효과적으로 학습시키며, inference 단계에서는 학습된 모델을 사용하여 노이즈화된 데이터를 복원한다.

1. 손실 함수(Loss Function) 정의
모델을 학습시키기 위해 negative sampling 기법을 활용한 손실 함수를 정의한다. 이 손실 함수는 사실 임베딩 간의 거리 계산을 기반으로 하며, 다음과 같이 표현된다:

$$L = -\log \sigma(\gamma - d_1(X_\tau, \text{Denoise}(X_\tau))) - \sum_{i=1}^n \frac{1}{k} \log \sigma(d_1(X_\tau, \text{Denoise}(X_{\tau_i})) - \gamma)$$

여기서 \(\gamma\)는 고정된 마진, \(\sigma\)는 시그모이드 함수, \(d_1\)은 \(L1\) 거리, \(X_{\tau_i}\)는 negative sample을 나타낸다.

2. Denoising 과정
모델이 예측한 노이즈와 최종 복원된 결과는 다음과 같은 수식을 통해 상호 변환될 수 있다:

$$\text{Denoise}(X_\tau) = \frac{1}{\sqrt{\alpha_t}} X_{\tau_t} - \sqrt{\frac{1 - \alpha_t}{\alpha_t}} \epsilon_\theta(X_{\tau_t}, t, X_h, X_r)$$

여기서 \(\epsilon_\theta\)는 모델이 예측한 노이즈 값이다.

3. Inference 과정
Inference 단계에서는 학습된 Conditional Fact Denoiser와 조건 임베딩(즉, \(X_h\)와 \(X_r\))을 사용하여, 노이즈화된 데이터를 원래의 상태로 복원한다. 이후, 복원된 사실 임베딩 \(X_\tau\)를 기반으로 타겟 엔티티 \(X_t\)를 예측하고, 각 엔티티 간의 거리를 계산하여 최종 예측값을 도출한다.

요약하면, Training 단계에서는 negative sampling을 통해 모델을 학습시키고, Inference 단계에서는 학습된 모델을 사용하여 노이즈화된 데이터를 복원하여 타겟 엔티티를 예측하는 과정을 수행한다.

Experiments

Main Result - Knowledge Graph Completion(WN18RR, FB15k-237)

FB15k-237에서 특히 높은 성능을 보여주었다. SOTA달성.

Main Result - Knowledge Graph Completion(Kinship, UMLS)

Main Result - Knowledge Graph Completion by relation types

FDM 모델은 모든 관계 유형(1-1, 1-N, N-1, N-N)에서 기존의 다른 모델들보다 높은 성능을 기록하고 있다. 특히 복잡한 관계 유형(1-N, N-1, N-N)에서 FDM은 우수한 성능을 보이며, 이는 FDM이 다양한 관계 패턴을 더 잘 모델링하고 처리할 수 있음을 보여준다. 반면, 비교적 간단한 1-1 관계에서는 다른 모델들과 비슷한 성능을 보이지만, 여전히 전체적으로 더 나은 결과를 나타낸다.

이 결과는 FDM이 KG에서의 다양한 관계 유형을 효과적으로 학습하고, 복잡한 패턴에서도 높은 정확도를 유지할 수 있는 강력한 모델임을 시사한다.

Limitations and Contributions

• Limitations
  • 조건부 사실 노이즈 제거기 설계: FDM은 노이즈화된 사실 임베딩을 원래의 타당한 상태로 복원하는 “Conditional Fact Denoiser”를 제안한다. 이 모델은 복잡한 KG 패턴을 효과적으로 학습하고, 다양한 관계 패턴을 포괄할 수 있는 구조를 설계했다.
  • MLP 기반의 효율적인 구조 사용: 이 논문은 KG의 특성을 고려하여, 복잡한 모델 대신 MLP 기반의 간단하고 효율적인 구조를 채택했다. 이를 통해 모델의 복잡성을 줄이면서도, 우수한 성능을 달성할 수 있었다.
• Contributions
  • Diffusion 모델의 계산 복잡도: FDM은 디퓨전 모델을 기반으로 하고 있어, 계산 비용이 높은 편이다. 이는 특히 대규모 KG나 실시간 애플리케이션에서 사용하기에 제약이 될 수 있다.
  • KG의 불완전성에 대한 민감도: FDM은 KG의 기존 데이터로부터 타당한 사실의 분포를 학습하므로, 학습 데이터가 불완전하거나 편향되어 있을 경우 성능이 저하될 수 있다. 이는 실제로 KG가 완벽하지 않은 경우, 모델의 예측력이 떨어질 가능성을 내포한다.
  • Large-Scale KG에 대한 성능 결과 실험 부재: Wikidata5M이나 YAGO3-10과 같이 규모가 큰 데이터셋에 대해서도 실험을 진행했어야 한다. 이는 모델의 일반화 능력과 관계있기 때문이다.

GR 카테고리 내 다른 글 보러가기