[2020] ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT

Document 랭킹을 성능 및 서비스 측면에서 고민한 ColBERT 논문에 대해 정리하고자 한다. 

 

1. Introduction

BERT 등의 큰 언어모델(LM)은 성능이 우수하나 계산량이 많아 latency가 생긴다. 때문에 query-document pair를 입력을 받아 수십만개의 문서 중 정답이 있을만한 문서를 찾는 문서탐색의 경우 시간이 매우 오래걸린다.

ColBERT는 이를 해결하기 위한 방법으로 1) 미리 계산한 문서 임베딩을 저장해두고(pre-computation) 2) query를 입력으로 받아 계산량이 낮은 연산(late interaction)으로 문서탐색을 수행하는 모델이다.

 

 

2. ColBERT 모델

1) Architecture

ColBERT는 두 개의 encoder와 late interaction을 위한 MaxSim 연산으로 구성되어있다.

 

2) Encoder

$f_{Q}$는 query $q$를 입력으로 받아 token별 임베딩 $E_{q}$을 계산하는 인코더, $f_{D}$는 document $d$에 대한 token별 임베딩 $E_{d}$을 계산하는 인코더이다. 각각 BERT 결과에 CNN를 붙여 차원을 축소시킨 임베딩을 얻는다. $f_{Q}$는 실시간으로 들어오는 query를 계산하고 $f_{D}$는 문서임베딩을 미리 계산해서 저장해둔다.

$$\begin{array}{lcl}
E_{q} = \text{Normalize} (\text{CNN}(\text{BERT}( \ [Q]q_{0}q_{1}...q_{l}\text{[pad]}...\text{[pad]}] \ ]))) \\
E_{d} = \text{Filter}(\text{Normalize} (\text{CNN}(\text{BERT}( \ [D]d_{0}d_{1}...d_{n}] \ ]))))

\end{array}$$

 

3) Score 계산(Late Interaction)

ColBERT에서는 $E_{q}, \ E_{d}$의 relevance를 계산하기 위해 MaxSIM 연산으로 제안하였다. 아래 수식처럼 $E_{q}, \ E_{d}$ 각각 token별 dot product의 max값을 sum하면 낮은 연산량으로 준수한 수준의 relevance score를 계산할 수 있다.

$$S_{q,d} = \sum_{i \in E_{q}} \underset{j \in E_{d}}{max} \ E_{q_{i}} \cdot E_{q_{j}}^{T}$$

 

 

3. 결과

ColBERT는 retrieve, reranking에서 BERT base와 비슷한 성능을 보이면서 계산량은 150배가 넘게 줄어든 모습을 보여준다.

 

또한 학습시나 추론시나 모두 실제 서비스에서 사용 가능할만한 수준의 latency를 보이고 있다.