원제: ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT

1. Introduction

정보 검색은 수작업 -> 임베딩 변환 후 내용 간 관계 모델링 -> LM(BERT)을 순위화 작업에 파인튜닝으로 발전해왔지만, 계산 비용이 100배에서 1000배까지 증가한다는 단점이 존재했다.
이를 해결하기 위해 BM25와 같은 전통적인 모델을 섞어서 사용하기도 하지만, 이 경우 검색 정확도가 크게 감소한다.
이를 해결하기 위해 BERT 기반 문맥화된 late interaction에 기반한 ranking model이 ColBERT이다.

late interaction: Query와 Document를 각각 두 개의 문맥화된 embedding 집합으로 독립적으로 인코딩하고, 이 두 집합 사이의 계산을 통해 관련성을 평가한다.

상호작용 중심 모델은 IR에서 성능이 좋고, 표현 중심 모델은 Query와 Document의 게산을 분리하면서 Document쪽을 오프라인에서 계산할 수 있게 해주어 Query당 계산 비용이 크게 줄어든다. ColBERT에선 두 장점을 합치는 것이 목적이다.

상호작용 중심 모델: Query의 token과 Document의 token을 하나하나 직접 비교하고 결과를 모아서 점수를 계산
표현 중심 모델: Query와 Document를 각각 하나의 벡터로 요약해서 유사도를 게산

2. Related Work

Neural Matching Models

KNRM: 상호작용 행렬로부터 matching signal을 추출하기 위해서 미분 가능한 kernel-polling techinique를 이용한다.

kernel-polling: query-document간 유사도 행렬을 kernel이라는 구간으로 나눠서 사용. 즉, 비슷한 단어들도 점수에 반영

Duet: exact-match-based와 embedding-based 유사도를 결합해서 순위를 매긴다.

Exact-match: BM25와 같이 질의 단어가 그대로 등장하는가
Embedding: 단어는 다르지만 의미가 비슷한가

ConvKNRM: Query-document간 n-gram을 매칭한다

n-gram: 단어 하나가 아닌 단어 덩어리를 맞춰본다. 이때 kernel-pooling으로 처리한다.

fastTest+ConvKNRM: sub-word token embedding을 활용해 rare word문제를 해결한다.

단어를 n-gram으로 쪼개서 단어 벡터를 구성한다.

SNRM: Query와 Document를 한 latent term로 인코딩하고 inverted index를 활용해 문서를 표현할 수 있다. 이를 통해 end-to-end 검색이 가능해진다. 하지만 최신 기법들에 비해서는 성능이 뒤쳐진다.

Language Model Pretraining

BERT와 같이 비지도 학습으로 pre-training을 진행하고, 이를 downstream 작업에 fine-tuning한다. 일반적으로 Query-Document 쌍을 BERT에 입력하고, [CLS]토큰에 MLP를 얹어서 점수를 계산하는 방식이다.

BERT Optimizations

LM기반 Ranking model들은 계산 비용이 매우 높다. 이를 줄이기 위해 distilling, compressing, pruning 등 연구가 있지만, 대부분 범용적인 BERT를 위한 것이기에 IR에서는 제한적인 향상만 보인다.

Efficient NLU-based Models

doc2query처럼 비용이 큰 NLU연산을 오프라인에서 수행하는 방향이다.

3. ColBERT

Query-Document간 상호작용을 지연시키는 것으로 비용을 감소시키고 , 벡터 유사도 검색을 통한 end-to-end 검색을 지원할 수 있다.

Architecture

구성 요소
1) Query Encoder $f_Q$
2) Document Encoder $f_D$
3) Late interaction

Query q와 Document d가 주어지면 $f_Q$는 q를 fixed dimension 임베딩 집합인 $E_q$로 인코딩하고, $f_D$는 d를 또 다른 임베딩 집합인 $E_d$로 인코딩한다. 이때 두 임베딩은 다른 토큰들에 의해 contextualized 된다.
이후 $E_q$와 $E_d$를 바탕으로 late interaction을 통해 Query와 Document간 score를 계산한다. 여기서 late interaction은 MaxSim의 합이다.

Maxsim: 가장 큰 값 한개만 남기는 방법
여기선 $E_q$에 속한 각 벡터 v에 대해 $E_d$의 벡터들과 cosine 유사도 중 최대값을 구하고, 이 값들을 모두 합산하는 방식이다. (suqared L2를 이용하기도 한다)

이런 토큰별 매칭 점수를 합산해서 문제의 전체 관련성을 추정한다.
MaxSim이기에 계산 비용이 매우 낮고, top-k개의 검색을 위한 pruning이 가능하다.

Query & Document Encoder

Query와 Document 모두 하나의 BERT 모델을 공유하고, 특수토큰 [Q]와 [D]를 통해서 구분한다.

Query Encoder

$E_q:= Normalize(CNN(BERT("[Q]q_0q_1...q_l\#\#...\#")))$

Query q가 주어지면 BERT의 WordPiece tokenizer를 사용해 $q_1, q_2,...,q_l$ 형태의 토큰 시퀀스로 변환한다. 이후에 앞에는 [Q]토큰을 추가한다.([CLS] 바로 뒤)
사전에 정의된 길이에 비해 Query token이 작은 경우 [MASK]토큰(#)을 사용해 길이가 $N_q$가 되도록 padding하고, 길이가 초과할 경우 처음 $N_q$개만 남긴다. 이를 query augmentation이라고 부른다. 이를 통해 Query-based emgedding을 생성하도록 유도하며, 새로운 용어를 추가하거나 기존 용어의 중요도를 재조정하는 메커니즘으로 작동한다.
BERT가 생성한 각 토큰의 문맥화된 표현은 activation function이 없는 Linear layer을 통과하면서 출력 차원 m의 임베딩을 생성한다. 이는 Query encoder의 효율에는 관련이 없지만, document embedding의 저장 공간에 영향을 주고, GPU로 전송하는 비용에도 영향을 주어 Query를 처리하는 시간에도 영향을 끼친다.
마지막으로 L2 norm이 1이 되도록 정규화된다.

Document Encoder

$E_d := Filter(Normalize(Linear(BERT("[D]d_0d_1...d_n"))))$

$d_1, d_2,...,d_n$ 형태의 토큰 시퀀스로 분할되며 맨 앞에는 [CLS], 그 뒤에는 [D]를 추가한다. 단, [MASK]는 추가하지 않는다.
해당 입력은 BERT, Linear layer을 통과하고, 구두점(제거해야하는 문장부호 등)에 해당하는 토큰 임베딩을 제거한다.

Late Interaction

$S_{q,d} := \displaystyle\sum_{i\in||E_q||} \displaystyle \max_{j \in ||E_d||} E_{q_i} \cdot E_{d_j}$

Query embedding $E_q$에 포함된 각 벡터에 대해 Document embedding $E_d$에 포함된 모든 벡터들과 유사도 중 최대값을 구하고, 이를 Query token전체에 대해 합산한다. 이때 보통 유사도 척도로는 cosine 또는 quarted L2 distance를 사용한다.
ColBERT는 미분 가능하고 end-to-end로 학습된다. 또한 BERT는 추가로 도입된 선형 계층이나 특수 토큰 [Q], [D]의 임베딩도 학습한다. 이때 late interaction자체는 학습 가능한 파라미터가 존재하지 않는다.
만약 Query q, positive document $d^+$, negative document $d^-$로 구성된 <q, $d^+$, $d^-$>가 주어지면 각 문서에 대해 독립적으로 점수를 계산하고 $d^+$와 $d^-$에 대해 pairwise softmax cross-entropy loss를 사용해 학습한다.

Offline Indexing

Query와 Document 계산을 분리시키며 Document를 오프라인에서 계산할 수 있다.
문서 컬렉션을 batch 단위로 순회하면서 각 batch에 대해 document encoder $f_D$를 실행하고, 이런 임베딩을 저장한다.
이를 최적화하기 위해 GPU를 활용해 문서 batch를 병렬로 인코딩하고, batch 처리시에는 가장 긴 문서 길이의 문서에 맞춰서 모든 문서를 패딩한다.
WordPiece tokenization에서 많은 시간이 소비되지만 이런 전처리는 문서들 간에 독립적이기에 이를 CPU 코어에서 병렬로 수행한다.

Top-k Re-ranking

비교적 작은 수 k(1000)에선 각 후보 문서를 전부 다 평가한다.
Query serving sub-system은 사전에 계산되어 저장된 document embedding을 메모리로 불러오고, 각 문서는 임베딩 형태로 표현된다
Query q에 대해 Query embedding $E_q$를 계산하고, 동시에 k개의 문서에 대한 임베딩을 모아 3차원 텐서D를 구성한다. 이후 D는 GPU로 이동한다.
GPU에서는 $E_q$와 D 사이의 batch dot-product 연산을 수행하고(이때 min-batch로 나누어질수도 있다) 결과적으로 query와 각 document간 교차 매칭 행렬로 구성된 3차원 텐서가 생성된다. 이후 document token 차원에 대해 합산을 하고, k개의 문서를 이 점수에 따라서 정렬한다.
$l = |q|$인 단일 시퀀스만 BERT에 입력하기에 계산이 매우 저비용이고, k가 증가할수록 더 잘 확장이된다.

End-to-End Top-k

후보가 너무 많아 전부 다 평가하기엔 문제가 존재하는 경우, MaxSim의 Pruning에 유리하다는 점을 이용한다. Query 임베딩 하나와 Document 모두에 대해 Maxsim을 계산하는 대신, fast vector-similarity data structure를 이용한다. 여기선 FAISS를 사용한다.

Document

Offline indexing 이후에 각 임베딩과 document간 매핑을 유지하고 FAISS index에 삽입한다. 이때 문서별이 아닌 전체 컬렉션의 모든 토큰 임베딩을 한 바구니에 모은다.(V = {$v_1,v_2...,v_M$}

Query

ColBERT 점수 계산 방식을 기반으로 한 두 단계로 처리한다.
각 단계를 조금 더 나눠서 보면
1) 근사치 뽑아내기

• a) Query q를 $f_Q$로 처리해서 $E_q$ = {$q_1, q_2, ...,q_{N_q}$}
• b) 전체 V에서 $q_i$와 가장 유사한 임베딩 상위 k'(=k/2)개를 가져온다. 즉, Query token = 32개일때 $N_q \times k'$개의 document token을 가져온다.
• c) 해당 token들을 document로 변환한다.
  이때 document 개수 $K \le N_q \times k'$

2) 정밀

• 후보 문서 K개에 대해서 ColBERT대로
  $S_{q, d} = \displaystyle\sum_i \displaystyle\max_j sim(q_i, d_j)$를 계산해서 최종 top-k개를 만든다.

4. Experiment

BERT급 성능을 내면서 동시에 BERT보다 170배 이상 빠르고, FLOPs는 14000배 감소한 결과가 나온다.
또한 End-to-End Retrieval에서도 기존의 BM25, doc2query보다 성능이 우수하며 동시에 BM25없이도 직접 검색이 가능했다.
Ablation Study결과 하나라도 빠질경우 성능이 떨어지는 결과가 나왔고, Average similarity보단 MaxSim의 성능이 더 좋았다.

5. Conclusion

기존의 BERT를 채용하되, Query와 Document를 Late Interaction을 통해 분리시켰다. 이를 통해 Offline Indexing이 가능해졌다. 동시에 GPU 연산에 더 유리한 연산들만 남기면서 속도도 더 높인 구조이다.
이때 Dense과 달리 단어 토큰 자체는 그대로 유지한다. 때문에 용량을 많이 차지한다는 단점이 존재한다.