[2023] Deep Bidirectional Language-Knowledge Graph Pretraining

[2023] Deep Bidirectional Language-Knowledge Graph Pretraining

Michihiro Yasunaga, Antoine Bosselut, Hongyu Ren, Xikun Zhang, Christopher D Manning, Percy Liang, Jure Leskovec

 

본문의 논문은 aaai 2023 workshop paper로, 링크를 확인해 주세요.

 

 

 

Abstract

 

 저자들의 제안 모델인 DRAGON (Deep Bidirectional Language-Knowledge Graph Pretraining) 은 MLM 과 KG의 linke prediction, 두 가지의 self-supervised reasoning task를 통합한 언어-지식 설립 모델을 사전훈련한 방식을 사용하였다. text의 segment와 연관된 KG의 subgraph의 쌍을 input으로 넣고 두 modalities의 정보를 양방향으로 융합한다.

 이는 QA 와 생물학 도메인과 같은 다양한 downstream tasks에서 현존하는 LM과 LM+KG 모델의 결과 값보다 평균 5% 이상의 좋은 성능을 보여준다.

 

 

 Introduction

 

 

 DRAGON 모델은 두 가지 주요 요소로 이루어져 있다 : text 와 KG를 양방향으로 융합한 cross-modal model이며, text와 KG에 대한 공동 추론(joint reasoning)을 학습하는 self-supervised를 사용한다.

 위 그림을 보며 모델을 설명하자면, text corpus와 KG의 raw data를 가지고 모델에 대한 input을 만든다. 이 input은 corpus에서 text segment를 샘플링한 것과 entity linking을 통한 KG에서 연관된 subgraph를 추출하여 ($text, local $KG) 쌍을 얻게 된다. 저자들은 이 input을 융합된 표현으로 인코딩하기 위해 cross-modal 모델을 사용한다. 여기서 모델의 각 레이어는 LM과 GNN을 사용한 KG로 인코딩하였고, 이 둘을 양방향 양식의 상호 모듈(modality interaction module) 로 융합한 것이다.

 저자들은 이 모델을 두 가지 self-supervised reasoning tasks로 pretraining하는데, (1) input text에서 tokens를 mask 하고 예측하는 MLM 과 (2) input KG에서 edges를 drop하고 예측하는 link prediction을 거치게 된다. MLM은 모델이 text에서 masked tokens를 추론하기 위해 KG에서 구조화된 지식과 함께 text를 사용하고 (위 그림에서 "round brush" - "art supply" 는 KG에서 multi-hop path), link prediction은 모델이 KG에서 잃어버린 links를 추론하기 위해 텍스트 맥락에서 KG 구조를 사용한다 ("round brush could be used for hair" 임을 인지할 수 있음). 따라서 이 공동 목표는 KG 구조에 의해 텍스트를 기반으로 하고 KG는 텍스트에 의해 동시에 컨텍스트화될 수 있도록 하여 추론을 위해 텍스트와 KG 사이에 정보가 양방향으로 흐르는 심층적으로 통합된 언어 지식 pretraining 모델을 생성합니다.

 

 

 

 

 Approach

 

• Definitions

 

 text corpus $\mathcal{W}$를 text segment의 셋 $\mathcal{W} =${W}, 각 text segment는 tokens (words)의 sequence로 W = ($\mathcal{w}_1, ... , \mathcal{w}_I$)로 정의한다. 지식 그래프 multi-relational graph $\mathcal{G = (V, E)}$는 $\mathcal{V}$가 KG에서 entity nodes의 셋이고, $\mathcal{E} ⊆ \mathcal{V} \times \mathcal{R} \times \mathcal{V}$는 nodes와 연결된 edges (triplet)의 셋, $\mathcal{R}$은 relation type ${r}$의 셋으로 이루어져 있다. KG의 각 triplet $(h,r,t)$는 (Paris, in, France) 와 같은 사실로 표현할 수 있다. 아무런 전처리 없는 raw KG는 노드가 백만개씩 되는 큰 그래프가 될 수도 있는데, raw KG (local KG) 중 subgraph는 다음과 같이 정의합니다 : $G = (V, E)$ 이고, $V = \left\{v_1, ... ,v_J \right\} ⊆ \mathcal{V}, E ⊆ \mathcal{E}$. 저자들은 언어-지식 모델을 두 가지 함수의 복합체 $f_{head}(f_{enc} (X))$로 보는데, 여기서 $f_{enc}$는 input X = (text segment W, local KG $G$)를 가지고, 각 text token에 대해 맥락 벡터 표현 (H_1, ... , H_I)과 각 KG 노드 (V_1, ... ,V_J)를 만들어낸다. 언어 모델은 언어-지식 모델에서 KG가 없는 ($J$ = 0) 특별한 경우이고, head인 $f_{head}$는 사전훈련 단계에서 self-supervised tasks를 수행하기 위해 이런 표현을 사용하고 미세 조정 단계에서 downstream tasks를 수행한다.

 

 

1. Input Representation

 

 text corpus $\mathcal{W}$와 큰 지식 그래프 $\mathcal{G}$가 있을 때, (text segment W, local KG G) 쌍을 이용하여 모델의 input instance를 만든다. 이 각 쌍들은 대략적으로 의미론적으로 정렬되므로, text 와 KG 상호보완적인 추론을 학습할 수 있다. 특히나, $\mathcal{W}$에서 각 text segment W에서, 다음과 같은 KG retrieval 과정을 통해 $\mathcal{G}$에서 연관되는 local KG를 추출할 수 있다.

 

KG retrieval

 text segment W가 주어질 때, 노드의 초기 셋인 $V_{el}$를 얻기 위해 W에서 entity로 언급한 것을 $\mathcal{G}$의 entity nodes에 연결한다. 여기서 $\mathcal{G}$에서 전체 retrieved nodes V ⊆ $\mathcal{V}$를 얻기 위해 2-hop nodes를 추가한다. 마지막으로 E ⊆ $\mathcal{E}$를 얻기 위해 $\mathcal{G}$에서 이런 노드들에 걸쳐 있는 모든 edges를 추가한다. 여기서 마지막 local KG인 $G = (V,E)$뿐만 아니라 마지막 input instance $X = (W,G)$를 산출해 낸다. (이후, local KG를 KG로 언급할 예정)

 

Modality interaction token/node

 (text, KG) 쌍의 결과에서, text와 KG에 각각 special token (interaction token) $w_{int}, v_{int}$를 추가한다. 이는 저자들의 cross-modal 인코더에서 modality interaction을 위한 인터페이스 뿐만 아니라 각 modality에 정보 pooling point도 제공할 것이다. 특히나, $w_{int}$는 기존의 text $W =(w_1, ... , w_I)$에서 준비가 되고, 새로운 relation 정보인 $r_{el}$을 사용하여 기존의 KG $V_{el} ⊆ V = \left\{v_1, ... ,v_J \right\}$에 entity-linked nodes를 위해 $v_{int}$를 연결한다. interaction token과 노드는 전체 input의 pooled 표현을 만드는 데 사용되며, 예시로, 분류를 위한 미세조정 시에 사용한다.

 

 

2. Cross-modal encoder

 

 

 

 위 그림처럼 text 와 KG 사이의 상호작용을 위해, bidirectional sequence-graph 인코더를 사용한다. 이 인코더는 $f_{enc}$가 text의 토큰과 KG의 노드들을 얻고, 각 토큰과 노드의 혼합된 표현을 만들기 위한 다중 레이어에 정보를 서로 교환하기 위함이다.

 

$(\mathbf{H}_{int}, \mathbf{H}_1, ... , \mathbf{H}_I), (\mathbf{V}_{int}, \mathbf{V}_1, ... , \mathbf{V}_J) = f_{enc}((w_{int}, .w_1, ... , w_I), (v_{int}, v_1, ... , v_J))$     (1)

 

 

 저자들은 이런 인코더를 사용하는 동안, 존재하는 작업에 대한 통제된 비교를 위해 가장 성능이 좋은 sequence-graph 구조인 GreaseLM을 적용하였다. GreaseLM은 text-KG inputs를 합성하기 위해 Transformers와 GNN을 결합한 모델이다.

 GreaseLM은 input text를 초기 토큰 표현으로 매핑하기 위해 Transformer LM에서 N개의 레이어를 처음 사용하였고, input KG 노드들을 초기 노드 표현으로 매핑하기 위한 KG 노드 임베딩을 사용하였다.

 

$(\mathbf{H}_{int}^{(0)}, \mathbf{H}_1^{(0)}, ... , \mathbf{H}_I^{(0)})$ = LM-Layers $(w_{int}, w_1, ... , w_I)$               (2)

$(\mathbf{V}_{int}^{(0)}, \mathbf{V}_1^{(0)}, ... , \mathbf{V}_J^{(0)})$ = Node-Embedding $(v_{int}, v_1, ... , v_J)$     (3)

 

 

 이 후 이 토큰/노드 표현들을 함께 최종 토큰/노드 표현으로 인코딩하기 위해 text-KG 혼합 M개의 레이어를 사용한다.

 

$(\mathbf{H}_{int}, ... , \mathbf{H}_I), (\mathbf{V}_{int}, ... , \mathbf{V}_I)$ = Fusion-Layers ($(\mathbf{H}_{int}^{(0)}, \mathbf{H}_1^{(0)}, ... , \mathbf{H}_I^{(0)}) , (\mathbf{V}_{int}^{(0)}, \mathbf{V}_1^{(0)}, ... , \mathbf{V}_J^{(0)})$)     (4)

 

 각 fusion 레이어 $(\mathcal{l}$ = 1, ... ,$M$)는 다음과 같은 과정을 거친다. :

 

$(\tilde{\mathbf{H}}_{int}^{(\mathcal{l})}, \mathbf{H}_1^{(\mathcal{l})}, ... , \mathbf{H}_I^{(\mathcal{l})})$ = LM-Layer $(\mathbf{H}_{int}^{(\mathcal{l} -1)}, \mathbf{H}_1^{(\mathcal{l} -1)}, ... , \mathbf{H}_I^{(\mathcal{l} -1)})$               (5)

$(\tilde{\mathbf{V}}_{int}^{(\mathcal{l})}, \mathbf{V}_1^{(\mathcal{l})}, ... , \mathbf{V}_J^{(\mathcal{l})})$ = GNN-Layer $(\mathbf{V}_{int}^{(\mathcal{l} -1)}, \mathbf{V}_1^{(\mathcal{l} -1)}, ... , \mathbf{V}_J^{(\mathcal{l} -1)})$               (6)

$[\mathbf{H}_{int}^{(\mathcal{l})} ; \mathbf{V}_{int}^{(\mathcal{l})}]$ = MInt($[\tilde{\mathbf{H}}_{int}^{(\mathcal{l})} ; \tilde{\mathbf{V}}_{int}^{(\mathcal{l})}]$)                                    (7)

 

 

 여기서 GNN은 KG 노드의 그래프 구조 인식을 유도하고, [ · ; · ] 는 이어 붙이기, MInt (modality 상호작용 모듈)은 MLP를 통해 텍스트 쪽의 상호 토큰과 KG 쪽의 상호 노드의 정보를 교환한다. 

 

 

3. Pretraining objective

 

 DRAGON 모델은 text 와 KG에서 함께 추론하는 것을 학습하기 위해 사전훈련을 목적을 가진다. text 와 KG가 서로 정보를 알려주고, 이 모델이 양방향 정보 흐름을 학습하기 위해서, 두 가지의 self-supervised 추론 학습을 한다. 

 

 Masked language modeling (MLM)

 

 우리가 흔히 아는 BERT 와 RoBERTa와 같이, MLM은 non-masked context에서 masked tokens에 대한 추론을 위해 사전학습 된 언어모델이다.  특히, 저자들은 input에 text-KG 쌍을 넣었을 때, text 와 KG의 구조적 지식을 활용하여 학습할 수 있도록 기대한다. 앞서 본 예시와 같이, KG에서 "round brush" - "art supply"인 경로가 인식하는 것은 masked 토큰인 "art supplies"를 예측하는데 도움을 줄 수 있다.

 MLM 을 수행하기 위해 저자들은 input text에 [MASK] token을 포함한 토큰의 하위 집합 $M ⊆ W$을 마스킹하고, task head인 $f_{head}$가 원래의 tokens를 예측하기 위해 인코더에서 문맥에 맞는 토큰 벡터 {$\mathbf{H}_i$}를 가질 수 있게 선형 레이어를 사용한다. 그리고 목적 함수의 cross-entropy loss 는 다음과 같다. :

 

$\mathcal{L}_{MLM} = - \sum_{i \in M} log p(w_i | \mathbf{H}_i)$     (8)

 

 

 Link prediction (LinkPred)

 

 위의 MLM은 text 부분의 예측이었다면, link prediction은 input KG에 대한 예측으로 볼 수 있다. 예를 들어, "X 어머니의 남편은 Y이다." 라는 구성 경로를 가지고, "X의 아버지는 Y이다."로 missing link를 예측할 수 있다. 그리고 text 와 KG 쌍이 input으로 들어갈 때도, KG에서 missing link를 예측할 수 있는데, "round brush could be used for hair" 라는 text가 들어갔을 때, link prediction은 (round,_brush, at, hair)로 edge를 가지고 예측이 가능하다.

input KG에서 edge triplets의 하위 집합인 $S$ = {$(h,r,t)$} ⊆ $E$를 가질 때, task head인 $f_{head}$에서 KG 표현 학습 프레임워크를 적용한다. h와 t는 각각 entity 노드, r은 relation이며 각 벡터 표현은 $\mathbf{h,r,t}$이고, 모델의 긍/부정 triplets를 위한 스코어링 함수는 $\phi_{r} (\mathbf{h,t})$로 정의된다. 그리고 $\mathbf{h} = \mathbf{V}_h, \mathbf{t} = \mathbf{V}_t, \mathbf{r} = \mathbf{R}_r$, {$\mathbf{V}_j$}는 인코더에서 문맥적인 노드 벡터의 표현이며, $\mathbf{R} = \left\{\mathbf{r}_1, ... , \mathbf{r}_{|R|} \right\}$는 학습 가능한 관계 임베딩이다.  KG triplet 스코어링 함수는 다음 세 가지를 고려하였다. :

 

DistMult : $<\mathbf{h,r,t}>$, TransE : - || $\mathbf{h + r - t}$ ||, RotatE : - || $\mathbf{h ⊙ r - t}$ ||     (9)

 

 여기서 < ·, ·, ·  >는 선형 내적이고, ⊙는 Hadmard 곱이다. 가장 높은 $\phi$ 값은 ($h,r,t$)가 no edge 보다 edge, 즉 부정보다 긍정에 가까운 값을 가짐을 나타낸다. 그리고 최적의 목적 함수는 다음과 같다. :

 

$\mathcal{L}_{LinkPred} = \sum_{(h,r,t) \in S} \left(- log \sigma (\phi_r (\mathbf{h,t}) + \gamma) + \frac{1}{n} \sum_{\mathbf{(h', r, t')}} log \sigma(\phi_r (\mathbf{h', t'}) + \gamma) \right)$     (10)

 

 

 ($h',r,t'$)은 n개의 긍정 샘플인 ($h,r,t$)와 상응하는 부정 샘플들이고, $\gamma$는 margin, $\sigma$는 시그모이드 함수이다. 이 함수는 held-out edges $S$의 triplets가 긍정으로, 다른 랜덤 triplets는 부정으로 예측될 수 있게 하는 지표이다.

 

 

 Joint training

 

 DRAGON 모델을 위해 MLM 과 LinkPred가 같이 최적화하는 값을 가져야 하므로, $\mathcal{L} = \mathcal{L}_{MLM} + \mathcal{L}_{LinkPred}$의 손실 값을 가지게 된다. 

 

 

 

4. Finetuning

 

 마지막으로 DRAGON이 텍스트 분류와 multiple-choice QA(MCQA)와 같은 downstream task에어떻게 fintune하는 지 보여준다. MCQA와 같은 경우에 있는 질문과 답 선택을 이어 붙인 input text $W$가 주어질 때, 위의 2.1 과 2.2 과정을 거친다. 그 후 전체 input $\mathbf{X}$의 풀링한 표현을 계산한다 : $\mathbf{X} = $MLP($\mathbf{H}_{int}, \mathbf{V}_{int}, \mathbf{G}$). $\mathbf{G}$는 $\mathbf{H}_{int}$를 쿼리로 사용한 {$\mathbf{V}_j | v_j \in \left\{v_1, ... , v_J \right\}$}의 attention 기반 풀링을 나타낸다. 마지막으로 풀링된 표현 $\mathbf{X}$는 BERT와 RoBERTa와 같은 LM 에서 [CLS] 표현이 사용되는 같은 방식으로 downstream task를 수행하는데 사용된다. 

 GreaseLM과의 차이점을 보면, GreaseLM은 이 부분에서 설명한 것처럼 finetuning만을 수행 (LM finetuned with KGs) 하는 반면에, DRAGON은 self-supervised pretraining을 하여, LM pretrained + finetuned with KGs가 된다.

 

 

 

 

 Experiments

 

 

 

 

 

 Analysis

 

 

Effect of knowledge graph

 

 DRAGON의 첫 번째 주요 포인트는 KG와 함께 쓰임이다. 위에서 언급한 것처럼, multi-step 추론과 부정을 해결한 것과 같이 복잡한 추론과 성능 향상에 좋은 결과를 보이고 있다.

 

 Quantitative analysis

 

 위의 table 2 를 보았을 때, 복잡한 추론을 포함한 질문에서 downstream task에서 좋은 성능을 보임을 증명하고 있다. 복잡한 질문에 대한 여러 카테고리를 보여준다. :

 (1) 부정의 표현 (no, never)

 (2) 접속사 (and, but)

 (3) hedge (somtimes, maybe)

 (4) 전치사 단계

 (5) 언급되는 entity의 수

 

 위의 카테고리 중, (1)과 (2)은 논리적인 복합 단계의 추론, (4)와 (5)는 더 많은 추론의 단계 또는 제약들, (3)은 텍스트의 복잡한 뉘앙스를 나타낸다. 저자들의 모델은 baseline인 RoBERTa보다 전반적으로 훨씬 뛰어난 성능을 보이며, 이는 언어와 지식 그래프를 함께 pretraining이 추론의 성능을 끌어올린 것으로 본다. 또한 현존하는 KG-증강 QA 모델인 QAGNN과 GreaseLM보다도 더 좋은 성능을 보이고 있다. 이는 self-supervised로 더 크고 다양하게 pretraining한 것이 finetuning만 한 모델인 GreaseLM보다 더 일반적인 추론의 능력에서 좋음을 알 수 있었다.

 

 

 Qualitative analysis

 

 

 CSQA 데이터셋을 사용한 시각화 결과를 위와 같이 보이고 있다. (A1, A2)에서 entities의 전치사와 부정표현을 추가했을 때, (B1 → B2)에서는 기존 질문에 추가적인 문장을 추가했을 때 그에 맞는 attention weights를 추가하고 최종 예측을 하고 있다.하지만 RoBERTa와 GreaseLM에서는 정확한 답을 예측하는 데에 방해가 되는 요소이다. 이런 부분을 보았을 때, 더 복잡한 질문이 들어오면 DRAGON만큼의 좋은 성능을 보일 수 없을 것이다.

 

 

 

Effect of pretraining

 

 이 모델의 또 다른 핵심은 pretraining이다. machine learning에서 주요한 세 가지, 데이터, 업무 복잡성, 모델의 수용 능력의 pretraining은 downstream task 데이터가 복잡성이나 모델의 수용 능력보다 작더라도 가능하게 만들어 준다.

 

 

 Downstream tasks with limited data

 

 위의 table 1을 보면 finetuning한 GreaseLM보다도 전반적으로 꽤나 좋은 성능을 보임을 알 수 있다. 

 그리고 finetuning data의 10%만을 사용한 상황에서도 좋은 결과를 보임을 증명하였다. 

 

 

 Complex downstream tasks

 

 table 1, tabe 2 그리고 figure 2에서도 보이다시피 좋은 성능과 복잡성에 대한 해결도 잘하는 것을 보인다.

 

 

 Increased model capacity

 

 

 저자들은 모델의 수용 능력을 보기 위해 text-KG 혼합 레이어의 수를 5에서 7로, 마찬가지로 GreaseLM에도 늘려 비교하였다. 수용 능력을 늘리는 것이 finetuning만한 모델에게서는 큰 이점이 없음을 보이며 pretraining 모델은 더 좋은 성능을 보여주었다.

 

 

 

Design choices of DRAGON

 

 

 Pretraining objective

 

 현 모델 design의 중요한 점은 함께 사용한 pretraining objective인 MLM + LinkPred이다. 이는 특히나 OBQA 데이터 셋에서 눈에 띄는 성능 향상을 보여주며 text 와 KG의 혼합 모델에서 양방향 self-supervised task가 추론에 좋은 성능을 보였다.

 

 

 Link prediction head choice

 

 

 KG 표현 학습은 연구에서도 활성화되는 영역이고, 다양한 KG triplet 스코어링 모델(DistMult, TransE, RotatE)이 제안되고 있다. 이를 활용하여 저자는 DRAGON의 link prediction head에 여러 가지 스코어링 모델을 사용하였다. DistMult가 약간의 edge를 가지고 있어 사용하였고, 이를 토대로 저자들의 모델의 일반성과 다른 KG 표현 학습 기술과도 사용될 수 있음을 시사한다.

 

 

 Cross-modal model

 

 

 또 다른 주요 포인트가 양방향 text-KG 혼합 레이어에 cross-modal 인코더를 사용한 것이다. 단순하게 text 와 KG 표현을 끝에 이어붙인 것 보다 좋은 성능을 냄을 보인다.

 

 

 KG structure

 

 

 마지막 키 포인트는 sequence-graph 인코더를 통한 KG의 그래프 구조와 link prediction을 활용한 것이다. 비교한 convert to sentence는 template를 활용한 문장을 기존 KG에 넣어 triplets를 변환한 것을 main text input에 넣고 vanilla MLM pretraining을 수행한 것이다. (여기서 vanilla model은 어떠한 다양성도 섞이지 않은 순수한 모델을 일컫는다.) 이 결과에서도 저자들이 사용한 graph 구조는 좋은 성능을 내는데 도움이 되었음을 보이고 있다.

 

 

 

 

Conclusion

 

 본 논문의 모델인 DRAGON은 텍스트와 지식 그래프에서 양방향 언어-지식 모델을 학습한 self-supervised pretraining 방식이다. 일반적인 도메인과 biomedical 도메인에서도 좋은 성능을 보였으며, 복잡한 추론, 예를 들어 긴 텍스트를 가진 질의 응답이나 여러 단계의 추론에서도 뛰어난 성능을 보이고 있다.