[논문 리딩] KG-BERT: BERT for Knowledge Graph Completion

KG-BERT: BERT for Knowledge Graph Completion

키워드	CompletionGraphKGKGCPLM
year	2019
저자	Liang Yao
Venue	ArXiv 2019
Memo	KG-BERT.
분류	연구
DONE
생성 일시	@2023년 11월 21일 오후 2:30
최종 편집 일시	@2023년 11월 22일 오후 12:19
Working

@article{Yao2019KGBERTBF,
  title={KG-BERT: BERT for Knowledge Graph Completion},
  author={Liang Yao and Chengsheng Mao and Yuan Luo},
  journal={ArXiv},
  year={2019},
  volume={abs/1909.03193},
  url={https://api.semanticscholar.org/CorpusID:202539519}
}

LLM과 KGC를 합친 프레임워크에 대한 논문을 읽어보았다. 이 논문은 그 논문의 이전 논문으로 KG-LLM을 제대로 이해하려면, 이 논문에 대한 이해가 필요하다고 느껴서 읽게 되었다.

이 논문은 KG-PLM 중에서도 BERT 모델을 PLM으로 사용해서 triple의 semantic 정보를 PLM에 fine-tuning 한다.

Introduction

이 논문은 entities와 relations, triples를 textual sequence로 다뤄서 knowledge graph completion (KGC) 문제를 sequence classification 문제로 전환한다. 그리고 이 sequences에 BERT를 fine-tuning 해서 triple이나 relation이 올바른지 예측한다.

Contribution

• KGC를 위한 language modeling 방법을 처음으로 제안함

• triple classification과 relation predictions, link prediction, 세가지 KGC task에서 SOTA 달성.

Related Work

KG Embedding

Model Only using Structural information

KG를 임베딩하는 방법은 크게 두가지로 나눌 수 있다.

1. Translation distance model : 거리에 기반한 scoring functions를 사용해서 KG를 임베딩한다.
   • TransE (Bordes et al. 2013)
   • TransH (Wang et al. 2014b)

2. Semantic matching models : 유사도 기반 scoring functions를 사용해서 KG를 임베딩한다.
   • RESCAL (Nickel, Tresp, and Kriegel 2011)
   • DistMult (Yang et al. 2015)

Model using Textual descriptions

위의 모델들은 구조적인 정보만 사용한다. 하지만 다른 종류의 외부 정보 (ex. entity types, logical rules, textual descriptions) 를 도입하면 성능을 더욱 높일 수 있다. 이런 한계점을 보완하기 위해 textual 정보를 사용하는 모델이 제안되었다.

이 모델들은 Entity들의 이름이나 descriptions를 외부의 말뭉치로 학습된 word embedding을 통해 entity의 representation을 얻어 KG embedding에 사용했다. 하지만 이 모델들은 동일한 textual representations를 사용하기 때문에, entity/relation descriptions에 포함된 단어가 triples에 따라 다른 중요도를 가질 수 있다는 사실을 간과하게 된다.

Text-enhanced KG embedding model

위의 문제를 해결하기 위해 제안되었다. 다른 triples에 포함된 relation에 다른 임베딩을 할당하는 방법이다.

• (Xu et al. 2017) : LSTM에 self-attention 매커니즘을 사용해서 주어진 relations마다 contextual text representation를 만들었다.

• (An et al. 2018) : 특정 triple specific relation mentions와 relation mention과 entity description 사이의 mutual attention mechanism을 활용해서 정확한 TEKE 방법을 제안했다.

이 방법이 triple들의 다양한 entities와 relations의 의미들을 다룰 수 있음에도 불구하고, 이 방법은 entity descriptions와 relation mentions, word co-occurrence with entities 만을 활용하기 때문에 거대하고 scale free한 텍스트 데이터의 구문적이고 의미적인 정보들을 충분히 사용하지 못한다.

KG-BERT

PLM을 통해 풍부한 언어 정보와 함께 Context-aware text embeddings를 학습할 수 있다.

LM Pre-training

PLM은 크게 두 카테고리로 나눌 수 있다.

1. Traditional word embedding methods : feature-based 방법으로 문맥과 독립된 단어 벡터를 학습한다.

2. Fine tuning approach : pre-trained 모델 아키텍쳐와 파라미터를 특정 NLP task를 수행하는 시작점으로 삼는다. Pre-trained 모델은 텍스트에서 풍부한 semantic pattern을 capture 할 수 있다.

최근에는 PLM을 사용해서 KG의 문맥을 포착하는데 사용하는 연구들이 제안되고 있다.

• (Wang, Kulkarni, and Wang 2018) : KG에 Random walk를 해서 생성한 entity-relation chains (문장들)에서 문맥적인 임베딩을 학습한다. KG의 초기 임베딩으로는 TransE 같은 KGE 모델을 사용했다.

• (Zhang et al. 2019) : KG의 중요한 entity들을 포함해서 BERT의 language representation을 강화했다.

• (Bosselut et al. 2019) : GPT를 사용해서 주어진 head 구문과 relation type에 대한 tail 구문을 생성한다.

본 논문은 기존의 방법들과 달리 entity와 relation의 이름이나 설명을 입력으로 사용해서 BERT를 fine-tuning 시켜서 주어진 triple이 얼마나 말이 되는지 (올바른 triple인지) 점수를 계산한다.

Method

Knowledge Graph BERT (KG-BERT)

KGC를 위해 BERT를 fine-tuning한 모델이다. Name/description word sequences를 fine-tuning에 사용하는 input sentence로 사용했다. BERT는 문장을 연속적인 텍스트나 단어 sequence로 사용한다. 그래서 triple의 개연성을 모델링 하기 위해, 이들은 $(h,r,t)$의 문장을 하나의 sequnce로 만들었다.

입력 sequence 구축

• 모든 입력 sequence의 첫번째 토큰은 항상 특별한 classification 토큰 [CLS]이다.

• Head/Tail entity 문장 : entity 이름이나 설명

• Relation 문장 : relation 이름

  예시)

  head - “Steven Paul Jobs was an American business magnate, entrepreneur and investor.” 또는 “Steve Jobs”

  relation - “founded” tail - “Apple Inc. is an American multinational technology company headquartered in Cupertino, California.” 또는 “Apple Inc.”

• 각 문장들은 [SEP] 토큰으로 나뉜다.

주어진 토큰에 대해 입력 representation은 해당하는 토큰과 세그먼트, position 임베딩을 더해서 만든다. [SEP]으로 나뉜 요소들은 다른 세그먼트 임베딩을 갖는데, Entity와 relation은 다른 세그먼트 임베딩을 갖는다. (head, tail entity는 동일한 세그먼트 임베딩을 공유한다.)

• Segment embedding : 문장과 문장을 이어주는 용도로 표현.

  BERT의 입력으로 연결된 문장이 들어왔을 때, 모델이 문장을 구별하기 위한 것이다.

  참고 : https://choice-life.tistory.com/25

• Position embedding : 토큰의 sequential한 의미를 보유하기 위해 표현

  다른 토큰이지만 같은 위치 $i \in \{1,2,3, \dots, 512\}$ 에 있다면 같은 position embedding을 갖는다.

Notation

• $E_i$ : token $i$ 가 갖는 input representation.

• $C \in \R^H$ : [CLS]의 final hidden vector. Triple의 점수를 계산할 때 aggregate sequence representation로 사용된다.

• $T_i \in \R^H$ : $i$ 번째 입력 토큰의 final hidden vector.

• $W \in \R^{2\times H}$ : Classification layer weights. Triple classification fine-tuning을 할 때 도입되는 새로운 파라미터.

• $\bold{s}_\tau = f(h,r,t) = sigmoid(CW^T)$ : Triple $(h,r,t)$ 에 대한 scoring function.

  $s_{\tau 0}, s_{\tau 1} \in [0,1]$ 로 구성된 벡터이고, $s_{\tau0} + s_{\tau1} = 1$.

• $\mathbb{D}^+, \mathbb{D}^-$ : Positive triple set, negative triple set.

Loss function

cross-entropy loss를 사용한다.

$$\mathcal{L} = - \sum_{\tau \in \mathbb{D}^+\cup\mathbb{D}^-} (y_\tau \log (s_{\tau0}) + (1-y_\tau)\log(s_{\tau1}))$$

$\bold{s}_\tau$ 을 예측 확률로, $y_\tau \in \{0,1\}$ (negative or positive) 를 레이블로 사용했다.

Negative Sampling

$\mathbb{D}^-$ 는 $(h,r,t)\in\mathbb{D}^+$ 의 head entity $h$나 tail entity $t$를 임의의 entity $h'$ 이나 $t'$ 로 바꾸는 방식으로 간단하게 생성했다. 만약 생성된 triple이 $\mathbb{D}^+$ 에 포함된다면 $\mathbb{D}^-$ 로 취급하지 않는다.

Weight Training

Pre-trained parameter 가중치와 새로운 가중치 $W$는 경사 하강법을 통해 업데이트 한다.

또 이 논문은 relation prediction task를 위해 다른 형태의 아키텍쳐의 KG-BERT를 제안했다.

두 entities $h$ 와 $t$의 문장만을 사용해서 둘 사이를 연결해주는 relation을 예측하는 아키텍쳐이다.

논문에 따르면, figure 2에 해당하는 아키텍쳐를 사용해서 곧장 relation을 예측하는 것이 figure 1에 해당하는 아키텍쳐로 relation corruption (relation $r$ 에 임의의 relation $r'$를 바꿔 넣어 negative triple을 생성하는 것) 을 통해 relation을 예측하는 것보다 낫다는 것을 알아냈다.

Figure 1로 수행한 relation prediction

• $W' \in \R^{R\times H}$ : Classification layer weights. Relation 예측 fine-tuning 단계에서 도입된다.

• $R$ : KG에 존재하는 relation의 종류 수.

• $\bold{s}'_\tau = f(h,r,t) = softmax(CW'^T)$ : Scoring functions.

  $\bold{s}'_\tau \in \R^R$. $s'_{\tau i} \in [0,1]$ 로 구성된 $R$-차원의 실수 벡터이고 $\sum^R_i s'_{\tau i} = 1$ 이다. 전체를 합하면 1이 나오므로 확률에 대한 term으로 이해할 수 있다. 이 벡터에서 $i$-번째 요소가 가장 높은 score를 갖는다면 relation $r_i$ 가 두 entity를 이어주는 relation 일 것이라고 예측한 것이다.

이때 사용한 loss :

$$\mathcal{L}' = -\sum_{\tau \in \mathbb{D}^+}\sum^R_{i=1} y'_{\tau i}\log(s'_{\tau i})$$

• $y'_{\tau i}$ : triple $\tau$ 를 위한 relation indicator.

  Triple $\tau = (h,r,t)$ 에 포함된 relation $r$ 에 대한 정보만 모으기 위한 term이다. $r=i$ 이면 $y'_{\tau i} = 1$ 이고, $r \neq i$ 이면 $y'_{\tau i}=0$ 이다.

Experiments

본 논문에서는 KGC의 3가지 task에 대해서 KG-BERT의 성능을 평가했다.

1. Triple classification : 주어진 Triple이 올바른 (KG에 타당한) triple인지를 예측.

2. Link (Entity) prediction : 하나의 entity와 relation이 주어졌을 때, 다른 entity를 예측.

3. Relation Prediction : 두 개의 entity가 주어졌을 때, 두 entity를 이어줄 수 있는 relation을 예측.

Dataset

Dataset	Data	Tasks	Entity input sentence	Relation input sentence
WN11	WordNet	Triple Classification	Entity names	Relation name
FB13	Freebase	Triple Classification	Entity description	Relation name
WN18RR	WordNet	Link/Relation prediction	Synsets definitions	Relation name
FB15K	Freebase	Link/Relation prediction	Entity names	Relation name
FB15k-237	Freebase	Link/Relation prediction	Entity description made by (Xie et al. 2016).	Relation name
UMLS	Medical semantic network	Link/Relation prediction	Entity names	Relation name

WordNet은 영어에 대한 어휘 KG 이다. 각 entity들은 동의어의 집합이다. Freebase는 일반적인 사실에 대한 거대 KG 이다. UMLS는 의학적 의미에 대한 네트워크로 의미적 타입 (entity) 들과 관계들을 포함한다.

WN11과 FB13은 test set에 positive, negative set을 가지고 있어서 triple classification에 사용할 수 있다. 반면에 나머지 데이터셋들은 test set에 positive 한 triple들만 갖고있기 때문에 link prediction과 relation prediction에 사용했다.

Entity에 대한 input sentence로는 Entity name, description이나 synset (동의어 집합) 정의를 사용했고, relation에 대한 input sentence로는 전부 relation name을 사용했다. 자세하게는 위의 표를 참조하면 된다.

Baselines

SOTA KG embedding 방식 : 구조적 정보만을 사용

• TransE

• TransH

• TransD

• TransR

• TransG

• TranSparse

• PTransE

• DistMult

• DistMult-HRS

Neural tensor network :

• NTN

• ProjE

CNN model :

• ConvKB

• ConvE

• R-GCN

Textual information을 사용한 KG embeddings :

• TEKE

• EKRL

• AATE

Entity hierarchical types를 사용한 KG embeddings : TKRL

Contextualized KG embeddings : DOLORES

Complex-valued KG embeddings :

• ComplEx

• RotatE

Adversarial learning framework : KBGAN

Settings

Pre-trained BERT-Base model (initialization of KG-BERT)

12 layers, 12 self-attention heads, $H=768$.

저자들은 선수 실험으로 BERT-Base 모델이 BERT-Large 모델보다 더 나은 결과를 보여준다는 것을 확인했다. 그리고 BERT-Base 모델이 Large 모델보다 더 단순하고 하이퍼 파라미터에 덜 민감하다는 사실 또한 알아냈다.

Fine-tuning :

• Adam

• Batch size : 32, learning rate: 5e-5, dropout rate: 0.1.

  다른 하이퍼 파라미터들도 실험해봤지만 많은 변화를 주지 않았다.

• Epoch :
  • Triple classification : 3
  • Link prediction : 5
  • Relation prediction : 20

  실험을 통해 relation prediction의 경우 epoch 횟수가 늘어날 수 록 성능이 올랐지만, 다른 두 task의 경우에는 해당되지 않았다는것을 알아냈다.

• Negative sampling
  • Triple classification : 1개의 positive triple에 대해 1개의 negative triple. 이진 분류를 할 때 밸런스를 맞추기 위함.
  • Link prediction : 1개의 positive triple에 대해 5개의 negative triples. 1, 3, 5, 10으로 실험을 해본 결과 5개씩 샘플링 한 경우의 성능이 가장 좋았음.

Results

Triple Classification

다른 베이스라인 모델에 비해 크게 향상된 결과를 보였다.

• TransE는 1-to-1 관계에 대해서만 다룰 수 있기 때문에 다른 Trans* 방법들 보다 낮은 성능을 보인다.

• DistMult는 상대적으로 괜찮은 성능을 보이지만, hierarchical 관계 구조 정보를 사용하는 DistMult-HRS에서 성능이 더 개선되었다.

• ConvKB의 경우, CNN 모델이 entity와 relation embedding 간의 global interactions를 잘 포착하기 때문에 괜찮은 성능을 보인다.

• DOLORES 는 ConvKB에 entity-relation random walk chain의 문맥적인 정보를 더한 것으로 성능이 개선되었다.

• NTN은 특히 FB13에서 높은 성능을 보였는데, 이는 모델이 훌륭한 것 뿐만 아니라 entity들을 word embedding과 함께 표현하는 것이 도움이 된다는 것을 보여준다.

• 다른 Text-enhanced KG embedding (TEKE, AATE) 모델들은 그들의 기본 모델들 (TransE, TransH)를 뛰어넘는 것을 보여주어, 외부의 텍스트 데이터의 이점을 증명했다. 하지만 그럼에도 풍부한 언어 패턴을 제대로 활용하지 못했다는 한계를 갖는다.

• KG-BERT는 FB13보다 WN11에서 더 성능이 높게 올랐는데, 이것은 WorldNet이 언어적 KG이기 때문에 PLM에 담긴 언어적 패턴들과 더 가깝기 때문이다.

Training 퍼센트에 따른 실험 결과를 나타낸 figure 3을 통해 KG-BERT는 제한된 training triples에서도 높은 성능을 가질 수 있다는 것을 알 수 있다.

논문에서는 10% training triples인 세팅에서도 text-enhanced model들보다 높은 성능을 보였다고 말한다. 하지만 주어진 도표에 포함된 베이스라인들은 전부 structure 정보만 사용하는 모델로 알고 있고, 이와 관련된 실험의 결과를 제공하지 않기 때문에 수치적인 비교가 불가능 했다.

아무튼, 이러한 실험들의 결과로 KG-BERT는 거대한 외부의 텍스트 데이터가 갖는 풍부한 언어적 패턴을 충분히 활용할 수 있어 KG의 sparseness를 잘 극복할 수 있다고 주장한다.

Triple Classification에서 KG-BERT의 성능이 좋은 이유 :

1. Input sequence가 entity와 relation word sequence를 포함함.

2. Triple classification은 next sentence prediction task와 유사.

   BERT를 pre-training 할 때 두 문장 사이의 관계를 포착하는 next prediction task를 통해 하기 때문에 pre-trained BERT가 triple classification 에 맞게 가중치가 설정되었기 때문.

3. 다른 triple에 속한 같은 토큰들이 다른 벡터를 가질 수 있어서 triple의 문맥적 정보를 담을 수 있음.

4. Self-attention 매커니즘을 통해 triple fact와 연결된 가장 중요한 단어를 찾아낼 수 있음.

Link Prediction

하나의 entity와 relation이 주어졌을 때, 남은 entity를 예측하는 task이다. 남은 entity 자리에 KG에 포함된 모든 entity $e\in \mathbb{E}$ 를 넣어보고 점수를 계산해서 내림차순으로 등수를 매겨 예측을 수행한다.

평가 지표 :

• Mean Rank (MR) : 정답 entity 의 순위를 평균낸 것

• Hits@10 : Rank를 매긴 리스트에서 정답 entity가 상위 10에 포함되는 확률

따라서 MR은 낮을수록, Hits@10은 높을수록 좋다.

위의 표에서 클래식한 베이스라인 모델들 (초록색 하이라이트) 는 OpenKE toolkit을 사용해서 결과를 구했고 다른 결과들은 해당 모델들의 논문에 나온 결과값들을 반영했다.

• Hits@10의 경우 KG-BERT가 베이스라인보다 더 낮은 성능을 보였다. 저자들은 이 결과를 KG-BERT가 entity와 relation 문장의 의미적 관련성 KG 구조정보가 분명하게 모델링 되지 않았기 때문이라고 분석했다.

• CNN 모델들 중에서 ConvE와 ConcKB는 R-GCN보다 높은 성능을 보였다. 논문에서는 이 결과만 언급하고 이유에 대한 분석은 없었다.

• ComplEx는 WN18RR과 FB15k-237에서 좋지 않은 성능을 보였지만 adversarial negative sampling을 이용한 KBGAN과 RotatE에서 향상되었다.

Relation Prediction

평가 지표 : MR, Hits@1 (accuarcy)

위의 Table 4는 figure 2에 해당하는 아키텍쳐로 수행한 실험의 결과이다.

• KG-BERT의 성능이 Hits@1에서 가장 높은 성능을 달성했다. 이는 KG-BERT가 BERT fine-tuning에서 sentence pair classification하는 것과 유사하기 떄문에 BERT pre-training에서 이점을 얻을 수 있었다고 분석했다.

• Text-enhanced models (DKRL, SSP) 는 구조적 정보만 사용한 TransE, TransH 를 능가했다.

• TKRL, PTransE는 위계적인 entity 카테고리와 확장된 경로 정보에 대해 잘 동작했다

• ProjE는 KGC를 ranking problem과 ranking score vector를 최적화하는 것으로 다뤄서 좋은 성능을 보였다.

Discussions

• 한계점 : BERT model의 비용이 비싸다. 특히 link prediction을 할 때 모든 entity에 대해 추론을 해야하기 때문에 상당한 시간 비용을 필요로 한다.

• 해결 : ConvE 같이 1-N scoring 방식을 도입하거나 더 가벼운 LM을 사용.