ColBERT: Efficient and Effective Passage Search viaContextualized Late Interaction over BERT(2020) 리뷰

최근에는 계층적, Agentic, Modular 등의 다양한 RAG 방법론이 나오는데 basic 한 접근방법은 100만 개 이내의 embedding에 대한 exhaustive search 이후 k=1000 정도의 후보군에 대한 rerank를 사용하는 것이다. 

흔히 사용되는 Reranker는 보통 Cohere reranker나 오픈소스 모델을 사용하는데 이들의 base가 되는 알고리즘이 ColBERT이다. RAG에서 DPR 정도의 위치에 있다고 보면 될 것 같다.

Abstract&Introduction

• 당시 IR에 language model(LM)을 사용하는 방법이 활발히 연구되고 있었지만 이 방법은 하나의 score를 얻기 위해 query, document쌍을 동시에 처리해야 하는 단점이 있다.
• 이러한 단점을 극복하기 위해 ColBERT라는 ranking model을 제안한다. ColBERT는 query와 document를 별개로 encoding하고 유사성을 얻는 late interaction을 사용한다.
• pruning friendly한 방법으로 벡터 유사도를 구하여 모든 document에 대해 빠르게 점수화가 가능하다

ranking model은 a->d로 순서대로 발전됐다고 볼 수 있는데 a->c로 갈수록 성능이 오르는건 자명하다. 하지만 그림 1에서 볼 수 있듯 성능 향상에 비해 들이는 시간, 자원이 기존 ranking model에 비해 비교할 수 없을정도로 커졌다.

 

저자는 BERT와 같은 deep language model들이 query와 document를 pair로 처리하면서 contextualization(문맥 이해도 정도로 이해하면 될 것같다)이 깊어졌다고 표현했다. 그리고 이러한 LM들의 efficiency와 contextulization을 reconcile시키기 위한 기법인 ColBERT를 제시했다.

COLBERT

As introduced earlier, delaying the query–document interaction can facilitate cheap neural re-ranking (i.e., through pre-computation) and even support practical end-to-end neural retrieval (i.e., through pruning via vector-similarity search).

 

저자의 생각이 그대로 담겨있는 문장인데 여기서 delaying은 c(기존 SOTA)와 달리 d(ColBERT)에서는 q와 d embedding 사이에 연결점이 없다는 걸 의미한다. 즉, ColBERT는 q, d의 independant interaction, scoring via pruning을 구현한 방법론이다.

Architecture

q와 d가 입력되면 f_Q는 q를 E_q로 f_D는 d를 E_d로 인코딩한다. 그리고 E_q와 E_d의 relevance score를 maximum similarity(MaxSim)연산을 통해 구한다.

Query&Document Encoders

아키텍처적 변화는 없으며 Query, Document 각각을 encoding하는데 변화가 있는데

Query: SentencePiece로 토큰화 한 후 일정 길이의 threshold보다 길면 자르고 짧으면 [pad]를 붙여준다. 이후 기존 BERT의 [CLS]부분을 [Q] 토큰으로 변경한다.

 

Document: 동일하게 적절히 토큰화하는데 길이 제한을 두지 않으며 [CLS]토큰은 [D] 토큰으로 변경한다. 또한 각 d_n이 인코딩 된 후 구두점이나 특수기호에 속하는 embedding을 filtering해준다.

수식으로 표현하면 아래와 같다

​Late interaction

ColBERT는 BERT와 약간의 추가 파라미터([Q], [D])로 구성되어있어 end-to-end differentiable하다. 또한 interaction mechanism은 위에 보이는 것처럼 단순한 연산에 불과하다 따라서 추가로 훈련할 parameter도 필요로 하지 않는다.

Offline Indexing: Computing&Storing Document Embeddings

By design, ColBERT isolates almost all of the computations betweenqueries and documents, largely to enable pre-computing documentrepresentations offline.

저자가 계속 강조한 바와 같이 ColBERT가 기존 BERT base IR model과 다른 점은 Q와 D를 별개로(isolate) embedding하고 처리한다는 점이다.

이 방식의 장점은 모든 문서에 대해 미리 embedding하고 이 값을 저장하여 서비스단에서 처리시간을 줄일 수 있다는 점이다. 

Top-k Re-ranking with ColBERT

query embedding과 precomputing 된 document와의 score의 연산결과를 줄이기 위해 document 마다 하나의 score로 만들어야 한다. 저자는 이 부분에서 평균을 내는 방법이 아니라 가장 점수가 높은 score하나만을 사용한다.

따라서 각 query token에 대해 문서 전반에서 하나의 score를 얻게 되는 셈이다.

이 그림을 한번 더 보면 쉽게 이해할 수 있다. 각 MaxSim이 Score에 대한 embedding maxpooling이라고 볼 수 있다.

End-to-end Top-k Retrieval with ColBERT

하지만 만약 Document가 매우 큰 상황에서는 어떻게 대처해야 할까? 이때까지 이 논문의 가장 큰 contribution은 document embedding의 precomputation에 있었다.

이를 통해 미리 embedding을 저장하고 있으면 ANN(approximately nearest neighbor)을 사용해 검색 범위를 줄일 수 있다.

---

논문을 읽으며 자연스럽게 SentenceBERT 생각이 들었는데 offline indexing과 embedding pre computing 그리고 이로인한 ANN 사용까지...둘이 거의 동시에 개제되었고 찾아보면 둘이 서로를 인용하고 있지 않다.

또한 이 논문에는 언급되지 않았는데 Document를 토큰화한걸 embedding하고 저장하면 용량이 매우 커지니 사용시 주의해야한다.

 

논문링크: https://arxiv.org/abs/2004.12832

논문리뷰 from JinaAI(pseudo code가 있어 이해가 쉽다): https://jina.ai/ko/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/