[논문리뷰]SimKGC: Simple Contrastive Knowledge Graph Completion with Pre-trained Language Models

Wang, L., Zhao, W., Wei, Z., & Jingming, L. (2022). SimKGC: Simple Contrastive Knowledge Graph Completion with Pre-trained Language Models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). https://doi.org/10.18653/v1/2022.acl-long.295

Problem Statement

1. Limitation of Graph Embedding
Graph Embedding은 엔티티와 릴레이션을 Text description같은 추가 정보를 사용하지않고 Triple의 구조 정보를 저차원 벡터로 mapping하여 그래프를 학습하게된다. Graph Embedding 방식을 채택한 모델들은 TransE, TransR, RotatE등이 있다. Graph Embedding의 문제점은 바로 Graph에 내제된 정보중 오직 구조 정보(Structural Information)만을 사용해 학습하고, 텍스트 정보(Texture Information)는 사용하지 못한다는 점이다. 이로인해 그래프의 특징을 온전하게 반영하지 못한다.

1. Limitation of Text-Based Method
Text-based Method는 Graph Embedding과 달리 사용가능한 text를 entity representation learning을 위해 통합하는 방식이다. 즉, Text 정보를 이용하여 학습을 진행한다. 직관적으로 추가적인 input정보를 활용할 수 있기 때문에 Graph Embedding방식보다는 좋을 것이라 예측되지만, 실제로는 훨신 긴 Inference time이 소요되고 심지어 성능이 더 뒤쳐지는 결과를 보였다.

Knowledge Graph에 있는 text description

저자는 이 Text-based method의 성능 저하 이유를 Inefficiency in contrastive learning으로 본다. 즉, 기존의 Contrastive learning은 KGC를 하는데 pre-trainied 모델에 적용하기에 부적합하다는 것이다.

Related Work

1. Knowledge Graph Completion
Triple을 (head, relation, tail)로 정의하고, 불완전한 Triple (head, relation, ?) 또는 (?, relation,tail)이 주어졌을때 missing entity에 해당하는 ?를 찾는 task이다.

1. Pre-trained Language Model
[논문리뷰]BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
[논문리뷰]GPT-1:Improving Langauge Understanding by Generative Pre-Training

1. Contrastive Learning

Contrastive Learning은 데이터 인스턴스 간의 유사하고 비유사한 특징을 학습하는 자기 감독 학습의 한 형태이다. 이는 CV, NLP, 그리고 Knowledge Graph 등 다양한 도메인에서 적용되었다. KGC의 맥락에서 Contrastive Learning은 인공적으로 손상된 것들(Negative Sample)과 비교하여 Knowledge Graph에서 실제 Triple (h, r, t)를 구별하는 작업에 초점을 둔다.

• Positive Sample: Knowledge Graph에서 존재하는 실제 Triple (h, r, t). 예를 들어, 동물에 관한 Knowledge Graph에서는 (고양이, is_a, 포유류)와 같은 Triple이 있을 수 있다.
• Negative Sample: Knowledge Graph에서 존재하지 않는 인공적 또는 가짜 Triple. Negative Sample은 일반적으로 Positive Triple를 손상시키는 방식으로 생성되며, 이는 head 또는 tail을 임의의 엔티티로 교체함으로써 이루어진다. 예를 들어, Negative Sample은 (고양이, is_a, 파충류)일 수 있다.

Contrastive Learning 모델은 이러한 Positive Sample과 Negative Sample을 구별하는 방법을 학습한다. 훈련 동안 모델은 Positive Sample과 Negative Sample의 score를 조합한 Loss를 제시받는다. 이는 Positive Triple에 높은 점수를, Negative Triple에 낮은 점수를 부여하려는 목표를 가지고 있다. 시간이 지남에 따라, 모델은 실제와 가짜 Triple를 효과적으로 구별할 수 있도록 Knowledge Graph의 엔티티와 관계의 임베딩을 학습해야한다.

Knowledge Graph Completion Task에 대한 Contrastive Learning의 목적은 모델의 Knowledge Graph 내 복잡하고 다중 관계 데이터에 대한 이해를 향상시키는 것이다. 이는 모델이 missing link를 추론하고 새로운 ling를 예측하는 능력을 향상시킬 수 있다.

[Deep Learning]Contrastive Learning(대조 학습)이란?

Method

1. Notation

Knowledge Graph를 \(\mathcal{G}\), True triple를 (\(h,r,t\))라고 정의할 때, 이 논문에서는 추가적으로 Inverse Triple (\(t, r^{-1}, h\))을 정의한다. 이 때, \(r^{-1}\)은 \(r\)의 역방향 릴레이션이다. 이렇게 정의를 함으로써 결국 tail에 대한 link prediction을 진행하면된다는 이점이 생긴다.

2. Model Architecture

1) Cross-encoder vs Bi-encoder

Cross-encoder
SimKGC는 BERT와는 다른 Bi-encoder architecture를 사용한다. BERT는 cross 인코더이며 Classification또는 Regression을 하기 위해서는 하나의 인코더구조에 결과물을 내놓기 위한 layer를 추가로 쌓게 된다. 그리고 인코더에 문장 Pair(혹은 Triple)을 넣어주고 모델에서 나온 [CLS] 토큰을(혹은 Mean Pooling/ Max Pooling을 거친 벡터를) 추가된 layer에 넣어주어 결과를 도출한다. Pre-trained 모델을 가져와서 fine-tuning을 진행을 해보면 이처럼 두 개의 문장이 한 번에 [SEP]토큰으로 나누어진 구조이다.

Cross-encoder 모델에서 나오는 결과물은 결국 임베딩이 아니라 문장 Pair가 들어갔을 때의 연산 결과이다. 즉, [CLS] 토큰을 구해준 후에 그것을 추가적인 layer의 input으로 넣어주기 때문에 연산에 상당히 오랜 시간이 걸릴 수 밖에 없다. 또한 Test 환경에서도 해당 구조를 사용하면 문장 pair가 들어갈 때마다 유사도를 전체 모델에서 구해주어야 하므로 매우 높은 연산이 필요하다. n개의 데이터에 대해서 자신을 제외한 n-1개와 비교를 하며 학습해야 하므로 총 \(\frac{n(n-1)}{2}\)의 연산이 필요하다.

또한 이는 문장 간의 유사도를 계산해주는 모델을 학습하는 것이기 때문에 문장 임베딩을 인코더가 학습했다라고 보기는 어렵다. 단일 문장이 인코더에 들어가서 그것을 의미적으로 적절한 위치에 임베딩시켜주는 것이 목표가 아니라 문장 유사도를 얼마나 잘 내보내는 모델이냐가 목표이기 때문이다.

Bi-encoder
Cross-encoder의 단점을 보완하기 위해서는 인코더의 추가적인 Layer를 쌓는 것이 아니라 다른 외부 구조를 활용하여 유사도를 구해주면 된다. 이를 Bi-encoder라고 한다. Bi-encoder에서는 A, B문장은 각각 서로 다른 인코더에 들어간다. 문장들은 각각 임베딩이 되며 그것들의 결과가 Score를 연산해주는 구조를 통해 연산이 된다. 이 구조를 이용하면 미리 학습된 문장들(A, B문장)은 임베딩이 완료되어 저장되게 된다. Test 환경에서 문장들이 C, D문장들이 주어지면 모델은 이미 저장된 임베딩을 바탕으로 빠르게 임베딩을 연산해주게 되고 이들에 대한 결과물을 외부적 구조로 attention을 통해 유사도만 구하게 된다. 이를 통해 매우 빠르고 각각의 역할이 명확한 구조가 형성이 된다. 하지만, Bi-encoder하나의 성능은 Cross-encoder에 비해 성능이 떨어진다. 왜냐하면 Cross-encoder의 초점 자체가 그 결과물 자체를 잘 내기 위해 구현된 모델이기 때문이다.

SimKGC with Bi-encoder
Bi-encoder 구조를 사용하면 위의 두가지 문제점을 보완할 수 있다. 두 개의 인코더를 사용해서 두 입력 문장의 임베딩을 각각 따로 학습한다. 이렇게 구한 두 임베딩의 유사도를 나중에 계산하는 방식으로, Cross-encoder보다 훨씬 효율적으로 학습을 진행 할 수 있다. SimKGC역시 Bi-encoder를 통해 Time complexity 측면에서 이점을 가져갔다.

Bi-encoder에서 두 개의 인코더를 각각 \(BERT_{hr}\)과 \(BERT_t\)로 정의한다. 먼저 \(BERT_{hr}\)이다. 이 인코더는 head의 relation-aware 임베딩을 계산한다. SimKGC는 Text description을 사용하는 대표적인 모델이다. 입력으로 head와 tail의 text description을 [SEP] 토큰과 함께 받아 마지막 layer의 hidden state를 얻는다. 특이한점은 [CLS] 토큰처럼 첫 번째 토큰의 hidden state를 직접적으로 이용하지 않고, L2 normalization과 함께 mean pooling을 진행해 relation-aware 임베딩 \(e_{hr}\)을 얻는다.

• Encoder 1. \(BERT_{hr}\)
  • get relation-aware embedding \(e_{hr}\)
  • Input: text descriptions of \(h\) + [SEP] + text descriptions of \(r\) ​ - instead of using the hidden state of the first token (ex.[CLS]), use mean pooling followed by \({L_2}\) normalization

두 번째 인코더는 \(BERT_t\)이다. tail의 text description을 입력으로 받아 \(e_t\) 임베딩을 얻는다. 마찬가지로 \(e_t\) \(L_2\) norm으로 nomalization 되어있다.

• Encoder 2. \(BERT_t\)
  • get tail entity enbedding \(e_t\)
  • Input: only consists of the textual description for entity t
  • use mean pooling followed by \({L_2}\) normalization

이렇게 2개의 인코더로 이루어진 Bi-encoder를 통해 \(L_2\)로 normalized 된 임베딩(\(e_{hr}, \; e_t\))을 구했으면 이를 통해 유사도(similarity)를 계산해야한다. 논문에서는 Cosine Similarity를 사용하였다.

이렇게 유사도까지 구했으면 마지막으로 Tail Entity Prediction (\(h,r, ?\))을 하기위해 유사도를 Maximization해 score값을 얻어내야한다.

3. Negative Sampling

Knowledge Graph Completion을 하기위한 Training dataset은 오직 Positive sample만을 포함한다. contrastive learning을 하기 위해서는 하나의 positive triplet (h,r,t)에 대하여 하나 혹은 여러개의 negative triples이 필요하다. 기존의 방식은 보통 head 또는 tail를 랜덤하게 선택하여 false negative entity로 대체하는 방식(선택된 엔티티를 다른 엔티티로 교체)을 사용한다. 본 논문에서는 training efficiency를 개선하기 위하여 3가지 종류의 negatives를 함께 사용한다.

1) In-Batch Negative(IB)

Visual representation learning에도 자주 쓰이는 방식이다. 같은 batch안의 엔티티들은 negative sample로 사용될 수 있다. 이러한 In-batch negative들은 bi-encoder모델에 엔티티 임베딩을 재사용할 수 있게 해주기 때문에 효율적이다.

2) Pre-Batch Negative(PB)

In-batch 방식의 단점은 negatives의 수가 batch사이즈에 한정된다는 것이다. 보통 negatives의 수를 늘리는 것이 Contrastive learning에 효과적이라고 알려져 있다. Pre-batch 방식은 이전 배치들에 있는 엔티티의 임베딩을 negatives로 사용한다. 보통 1-2개의 Pre-batches가 사용된다.

3) Self-Negative(SN)

negative의 수를 늘리는 것 외에도 정답 entity와 가까이 있어 구분하기 힘든 negative를 사용하여 모델의 분별력을 키우는 것도 Contrastive learning에 있어서 중요하다. 이러한 negative를 Hard Negative라 한다. tail entity prediction의 경우 head entity \(h\)가 hard negative로 작용할 수 있다.

IB, PB, SN을 각각 \(\mathcal{N_{IB}}, \; \mathcal{N_{PB}}, \; \mathcal{N_{SN}}\)로 표기한다. 학습 중에 false negative가 발생할 수 있다. 예를 들어, N-N관계에 있는 트리플의 경우 negative sampling을 위해 head를 바꿨는데도 정답일 수가 있는 경우다. 이런 경우는 SimKGC에서 Binary Mask를 통해 filtering한다. 이 모든 Negative를 합쳐서 \(\mathcal{N(h,r)}\)로 표기한다.

논문에서 batch size는 1024이고, 2개의 pre-batch를 사용했다. 따라서 \(\vert \mathcal{N_{IB}} \vert = 1024 - 1\), \(\vert \mathcal{N_{PB}} \vert = 2 \times 1024\) 이며 \(\vert \mathcal{N_{SN}} \vert = 1\)이다. 따라서 총 negative수는 \(\vert \mathcal{N(h,r)} \vert = 3072\) 이다.

4. Graph-Based Re-ranking

Knowledge graphs에서는 멀리있는 엔티티들보다 가까이 있는 엔티티들끼리 더 관련성이 있는 Spatial locality inductive bias가 존재한다. Text-based KGC method는 이러한 inductive bias를 완벽하게 포착해내지 못한다. 논문에서는 이 점을 보완하기 위해 head entity \(h\)의 k-hop 이웃 노드 \(\mathcal{E_k}(h)\)에 존재하는 후보군 Candidate tail entity \(t_i\)에 대하여 더 높은 점수를 부여하도록 한다.

5. Training & Inference

1) Training with Loss

SimKGC는 InfoNCE loss를 추가적인 additive margin \(\gamma\)와 함께 사용한다.

• InfoNCE loss with additive margine

additive margin \(\gamma\)는 0보다 크면( \(\gamma \; > \; 0\), 양수)이면 모델이 정답 Triple (\(h,r,t\))의 점수, Score를 증가하게 유도한다. Loss식에서 \(\phi (h,r,t)\) 후보 Triple의 Score function이며 앞서 정의했듯, 이 Score function \(\phi (h,r,t) \; = \; cos(e_{hr}, e_t) \; \in \; [-1,1]\)가 된다.

Temperature \(\tau\)는 negative의 상대적인 중요도를 조정하는 역할을 한다. \(\tau\)가 작을수록 hard negative를 더 강조하게 된다. 논문에서는 \(\tau\)를 hyperparameter로 사용하지 않고, Learnable하게 만드려고 re-parameterize를 해 \(log \frac{1}{\tau}\)으로 만든다.

2) Inference

추론 시, 가장 시간 소모가 큰 Time-Consuming COst는 BERT에서 엔티티 임베딩의 정방향 forward pass를 계산하는 부분으로 Cost는 \(O(\vert \mathcal{E} \vert)\)이다. Test Triple의 총 개수는 \(\vert \mathcal{T} \vert\)이다. 각각의 Triple (\(h,r,?\))와 (?, r^{-1}, h)에 대해 relation-aware head entity embedding을 계산하고 내적(dot product)을 이용해 모든 엔티티들과의 ranking score를 계산해야 한다.

결론적으로 Bi-Encoder를 이용하는 SimKGC는 총 \(\vert \mathcal{E} \vert \; + \; 2 \; \times \; \vert \mathcal{T} \vert\)의 BERT forward passes를 거친다. 반면 Cross-Encoder를 이용하는 original BERT의 경우 \(\vert \mathcal{E} \vert \; \times \; 2 \; \times \; \vert \mathcal{T} \vert\) 만큼의 forward pass를 하므로 시간도 역시 더 오래걸린다. Bi-encoder 모델들의 경우 엔티티 임베딩과 찾아야 할 top-k개의 엔티티들을 효율적으로 Faiss라는 빠른 유사도 탐색 툴을 이용해 pre-compute할 수 있다.

Experiment & Result

1. Data Set

Text description을 이용하기 위해 WN18RR과 FB15k-237의 경우 KG-BERT에서 제공된 dataset을 그대로 사용했으며, Wikidata의 경우 이미 dataset안에 포함되어 있다.

• Evaluation Metric
  • Hits@k, k = 1,3,10: top-k의 랭킹된 엔티티들 중 정답 엔티티의 비율
  • MRR: Test triple의 reciprocal rank의 평균, 정답 엔티티가 랭킹중 어느 위치에 있는지에 중점을 둠

2.1 Knowledge Graph Completion with Wikidata

• Transductive setting : all entities in the test set also appear in the training set
• Inductive setting : no entity overap between train and test set

Wikidata5M 데이터셋의 Transductive setting에서는 거의 모든 metric에서 세 개의 배치를 한 번에 사용했을 때 성능이 좋게 나옴. Wikidata5M dataset에서 self-negative를 사용했더니 in-batch만 사용했을 때보다 10%이상 성능이 향상된걸 볼 수 있다. 참고로 Self-negative는 주어진 head에 대해 모델이 단순하게 1차원적으로 예측하는 것을 방해한다. 이러한 이유로 hard negative와 self-negative가 연결된다.

또한 inductive KGC에서 transductive setting보다 text-based model에 더 의존적이다. 따라서 Transductive setting과 inductive setting간 성능이 차이가 많이 난다.

2.2 Knowledge Graph Completion with WN18RR & FB15K-237

FB15k-237 dataset의 경우 embedding-based method가 더 좋은 성능을 보이고 있다. 이 데이터셋의 경우 훨씬 밀집된 그래프이기 때문에 의미적 연관성보다 일반화된 추론 능력이 모델에 더 필요하다고 설명하고 있다. FB15k-237의 경우는 이용가능한 정보에 근거해 예측을 하기 어려운 링크가 많다.

WN18RR의 경우는 SimKGC가 압도적으로 좋지만, In-batch에서의 성능이 대체적으로 우수하다. 다시 말해, Batch size만 크게 해줘도 성능이 좋아지는 것을 확인할 수 있다. 기존의 KGC text model들은 batch size가 매우 작았기 때문에 Sparse한 dataset인 WN18RR의 경우 batch size를 늘려서 성능이 올라간다.

future work : ensemble embedding based method with text-based methods

Inference time의 경우 Wikidata5M-Trans dataset에서는 SimKGC가 ~4.6 million ebeddings을 계산하는데 ~40분이 걸렸다고 한다. KG-BERT가 같은 조건에서 3000시간이 걸렸던 것과 비교하면 훨씬 단축된 결과다. (추론 시간을 획기적으로 줄인 첫 논문은 아니지만 의의 있음. ex.ConvE,StAR)

3.1 Ablation 1: What makes SimKGC Excel?

1) Two major factor

기존의 Text-based KGC 모델들과 달리, SimKGC는 크게 두 가지 차이점을 보인다.

• SimKGC는 더 많은 negative를 사용했다.(using more negatives)
• SimKGC는 Margin-based ranking loss 대신 InfoNCE loss를 사용했다.(switching from margin-based ranking loss to InfoNCE loss)

논문에서는 SimKGC in-batch를 학습할 때 InfoNCE loss에 256개의 batch size를 사용했다. 그 결과 가장 좋은 것을 확인할 수 있다. negative 수를 동일하게 설정하고 margin-based ranking loss와 InfoLoss만 비교해보더라도 성능차이가 많이나는 것을 알 수 있다.

➜ InfoNCE loss와 negative 수가 성능에 영향을 주는 직접적인 factor이다.

저자들이 보기엔 loss function이 더 큰 기여를 한다고 말한다. InfoLoss에서 hard negative는 더 큰 gradient값을 만들어 내고 negative 수를 늘리는 것은 결론적으로 더 robust한 representation을 이끌어 낸다.(Hardness-aware property가 결국 contrastive loss의 성능 향상에 기여한다.)

2) Margin-τ loss

위의 표4를 보면, 맨 마지막에 Margin-τ 가 있다. 이는 기존의 margin loss를 조금 변형한 것이다.

왼쪽처럼 InfoNCE loss대신 margin-based ranking loss를 사용하면 성능이 떨어진다. 하지만 여기서 오른쪽과 같이 margin-τ loss를 사용하면 왼쪽식을 이용했을 때보다 성능이 좋아지는 것을 확인할 수 있다. 이처럼 Loss에 따라 성능 차이가 큰 것을 알 수 있기 때문에, SimKGC에서 Loss가 얼마나 큰 contribution인지 알 수 있다. margin-τ loss는 마찬가지로 모델이 hard negative에 좀 더 주의를 기울이게해 성능을 향상시킨다.

3) negative 수에 따른 성능 비교

negative 수를 늘리면 늘릴수록 성능이 좋아진다. 하지만, negative수를 너무 늘리면 GPU memory 사용량이 기하급수적으로 증가하므로 optimization에 대한 cost가 증가한다.

3.2 Ablation 2: Re-ranking

1) re-ranking

re-ranking의 목적은 Knowledge Graph의 topological information을 간단한 방법으로 통합하기 위함이다. 즉, 구조 정보를 이용하기 위함이다. Graph의 connectivity pattern은 spatial locality를 나타내는 경우 re-ranking이 순위를 다시 지정하는 것에 대해 도움이 된다.

re-ranking이 성능에 미치는 영향을 알아보기위해 Wikidata5M-Trans 데이터셋으로 실험을 진행하였다. re-ranking을 하는 경우가 성능이 약간 더 좋은 것을 알 수 있다.

2) Fine-grained analysis

모든 릴레이션을 4개의 카테고리로 분류하여 Table 7과 같이 나눈다. 그리고 모델이 어떤 관계에서 성능이 좋은지 비교하기 위해 패턴별 MRR을 비교하였다. 그 결과 n-1관계가 가장 좋은 성능을 보여줬다. 그도 그럴것이, 정답이 n개가 되면 정확한 답을 고르기 힘들어지므로 성능이 떨어질 수 밖에 없다. 즉, 많은 수의 plausible한 정답이 진짜 True answer을 헷갈리게 만든다.

3.3 Ablation 3: Human Evaluation

MRR같은 metric으로 평가를 진행할 시 위와 같은 이유로 성능이 낮게 측정되는 문제가 발생할 수 있다. 따라서 본 논문에서는 human evaluation을 추가로 실행하였다. 100개의 랜덤하게 선택된 잘못된 prediction에 대하여 사람이 평가를 진행한 결과다.

• 49%의 wrong prediction에 대해서 사람은 올바른 정답이라고 평가했음을 확인할 수 있다.
  • How to accurately measure the performance of KGC is also an interesting future research direction.

Contribution

• KGC에 Contrastive Learning을 효율적으로 적용하였다.
  • InfoNCE
  • 3 types of negatives
  • Bi-enocder
• SOTA 기록

Reference

Cross-encoder와 Bi-encoder (feat. SentenceBERT)

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