[논문리뷰]ColBERT

Background

역색인(Inverted Index) 검색 엔진

후보 문서들을 검색하는 일도 중요하지만, 최종적으로 사용자에게 보여지는 문서를 고르는 Ranking작없도 중요하다. HotpotQA, MuSiQue, 2WikiMultiHopQA와 같은 데이터셋들은 이미 후보 문서들은 정해져있기 때문에, Ranking 작업이 필수이다.

전통적인 ranking방법으로는 BM25를 많이 사용하고 있다. BM25는 TF-IDF와 비슷하지만, 문서의 길이를 고려한 랭킹함수이다. 어떤 텀에 대하여 해당 텀이 어떤 문에서 얼마나 많이 등장하였는지(많이 등장할수록 값이 커짐)와, 얼마나 많은 문서에 등장했는지(많은 문서에 등장할수록 값이 작아짐)를 고려한다. BM25는 역색인 검색 엔진에서 미리 계산해 저장할 수 있기 때문에 많은 문서를 다루고 성능이 중요한 검색엔지에서 많이 사용된다.

역색인(Inverted Index) 검색 엔진은 문서 검색 시스템에서 효율적으로 데이터를 검색하기 위해 사용되는 방식 중 하나이다. 일반적인 색인이 문서와 그 안에 포함된 단어의 위치를 기록하는 방식이라면, 역색인은 각 단어에 대해 그 단어가 포함된 모든 문서를 기록하는 방식이다. 이를 통해 사용자가 특정 단어를 검색할 때 해당 단어가 포함된 문서를 빠르게 찾아낼 수 있다.

역색인에서는 키(key)가 단어가 되고, 값(value)은 그 단어가 등장한 문서 리스트 또는 문서 내에서의 위치 정보가 됩니다. 예를 들어, 다음과 같은 세 문서가 있다고 가정해 보겠습니다.

• 문서 1: “고양이는 귀엽다.”
• 문서 2: “강아지는 충성스럽다.”
• 문서 3: “고양이와 강아지는 친구다.”

역색인 검색 엔진을 이용하여 “고양이”라는 단어를 검색할 경우, 역색인은 즉시 문서 1과 문서 3을 반환할 수 있다. 대표적인 역색인 검색 엔진에는 엘라스틱서치(Elasticsearch), 아파치 루씬(Apache Lucene), 그리고 솔라(Solr) 등이 있다.

BERT를 활용한 Neural Ranking

1. BERT에게 \(\text{[CLS] Query [SEP] Document [SEP]}\) 형식으로 시퀀스를 입력시킨다. 이때 input data로 \(\text{[CLS]}\)와 \(\text{[SEP]}\) 토큰을 반드시 포함해야 한다.
2. 모든 BERT 레이어에 대해 실행한다.
3. 마지막 \(\text{[CLS]}\) 토큰의 출력 임베딩에서 score를 추출한다.

BERT는 높은 MRR 성능을 보여주었지만, computational cost가 매우 높은 것을 볼 수 있다. 위의 figure에 따르면 BM25와 BERT-large는 거의 100배 이상 차이가 난다.

Representation Similarity

DPR(Dense Passage Retrieval)은 BERT기반의 랭킹 방식에서 computational cost를 줄이고 MRR을 높이고자 시도한 대표적인 연구이다. DPR은 각 문서와 쿼리를 768차원의 벡터로 인코딩하여 계산하였다.

Representation Similarity 모델은 쿼리와 문서 임베딩을 미리 계산하고, 쿼리 임베딩과 문서 임베딩간의 내적(dot-product), 혹은 코사인 유사도(cosine similarity)를 사용한다. 굳이 Reranking에서 사용처를 한정할 필요 없이 End-to-End-Retrieval에 사용 가능하다. (Faiss) 하지만 이렇게 하나의 벡터로 표현하다보니 coarse grain representation이 되어버려 텀 레벨 interaction이 없어졌다. 그래서 BERT보다 낮은 MRR을 보이는 단점이 있다.

corse grain이란 큰 단위 또는 전체적인 수준에서 표현을 다루는 방식을 말한다. 즉, DPR(Dense Passage Retrieval)과 같은 방식에서는 쿼리와 문서의 전체적인 의미를 한 벡터로 추상화해서 처리하게 됩니다. 이때 세부적인 텀 단위(단어 수준)의 상호작용을 무시하게 되는 것을 “coarse grain” 표현이라고 부른다.

Fine grain representation에서는 쿼리와 문서의 각 단어(토큰) 또는 구문 간의 관계를 세밀하게 분석한다. BERT 같은 모델은 단어 수준의 상호작용(단어 간의 관계, 문맥)을 모델링할 수 있어서, 세밀한 의미 차이를 잡아낼 수 있다.

반면, DPR처럼 하나의 벡터로 쿼리와 문서를 표현하는 방법은 각 문서와 쿼리의 전체적인 의미만 반영할 뿐, 개별 단어 간의 세부적인 상호작용을 반영하지 못한다. 예를 들어, “사과”라는 단어가 쿼리와 문서에서 등장하더라도, 해당 단어의 위치나 문맥에 따른 의미 차이는 제대로 고려되지 않는다.

Problem Statement

기존의 IR (Information Retrieval)모델들은 매우 높은 컴퓨팅 자원을 요구하였으며, coarse grain하다.

Method

Retrieve를 함에 있어서 효율적이고 정확한 검색을 위해 모델에 요구되는 특성은 크게 세 가지이다.

• 독립적인 인코딩 가능
• Fine grained representation
• Vector similarity search를 통해 End to End Retrieval 가능

ColBERT는 이를 Late Interaction을 통해 가능하게 하였다. Late Interaction은 **Representation Similarity 모델처럼 Document를 미리 offline에서 계산해 computation을 빠르게 만들었다.

Model Architecture

ColBERT의 아키텍처는 다음과 같이 크게 세 개의 모듈로 구성된다.

• Query Encoder
• Documents Encoder
• Late interaction

ColBERT는 Late Interaction을 사용해 쿼리와 문서를 각각 임베딩하지만 하나의 BERT 모델을 공유하기 때문에 토큰 Speical 토큰 \(\text{[Q]}\)와 \(\text{[D]}\)를 사용해 인풋에서 쿼리와 문서를 구분해야 한다.

Query Encoder
Query Encoder는 주어진 쿼리 q에 대해 BERT-based WordPiece 토큰 \(q_1, q_2, \dots, q_l\)로 토크나이즈 한다. 그리고 토크나이즈된 쿼리의 시퀀스 앞에 Special 토큰 \([Q]\)를 붙인다. 만약 쿼리가 미리 정해진 토큰 개수 \(N_q\)보다 적다면, BERT의 \(\text{[Mask]}\)토큰을 붙여 패딩한다. (\(\text{[Mask]}\) 토큰으로 패딩하는 것을 Query Argument라고 함.)

$$E_q = \text{Normalize}(\text{CNN}(\text{BERT}(\text{[Q]}, q_0, q_1, q_2, \dots, q_l)))$$

Linear layer는 인코더로부터 representation을 넘겨받아 m-차원 임베딩을 만든다. ColBERT의 임베딩 차원수는 쿼리 실행 시간과 문서의 공간 관리에 큰 영향을 미친다. 시스템 메로리에서 GPU로 문서 representation을 옮겨야하기 때문이다. ColBERT의 reranking에서 가장 비싼 스텝은 모으고, 쌓고, 이를 CPU에서 GPU로 옮기는 작업이다.

Document Encoder
Query Encoder와 유사한 구조를 가진다. 문서 \(d\)를 토큰 시퀀스 \([d_1, d_2, d_3, \cdots, d_m]\)으로 만들고 BERT의 시작 토큰인 \(\text{[CLS]}\)를 문서 입력을 알리는 Speical 토큰 \(\text{[D]}\) 앞에 붙인다.

$$\text{Filter}(\text{Normalize}(\text{CNN}(\text{BERT}(\text{[CLS][D]}, d_o, d_1, \dots, d_m))))$$

중요한 점은, 문서 임베딩을 구함에 있어, [MASK]토큰은 절대 사용하지 않는다. 이후 미리 정한 리스트를 통해 특수 문자등의 기호 등을 필터링하여 문서당 임베딩 개수를 줄이는 과정을 거친다.

Late Interaction
쿼리 \(q\)와 문서 \(d\)간의 연관성은 relevance score라 하고, 이는 \(S_{q, d}\)로 나타낸다. \(S_{q, d}\)는 코사인 유사도나 Squared \(L_2\) distance를 통해 구한 값들을 합쳐 계산한다.

Late Interaction의 의미는 다시 말해, 하나의 인코더로 쿼리와 문서를 임베딩하지만, 각각 독립적으로 임베딩해 attention을 계산하지 않고, 나중에 따로 유사도 연산을 진행함을 의미한다.

Offline Indexing
ColBERT는 문서와 쿼리 계산을 분리시켜 문서 representation을 미리 계산하는 것이 가능하다. 색인 생성 과정에서 배치로 컬렉션에 있는 문서들을 document encoder로 임베딩을 계산하고 이를 저장한다. 배치 작업시 배치 문서 중 가장 긴 문서 길이로 모든 문서에 패딩을 한다. 이때 문서들을 길이로 정렬하고 여기서 다시 배치로 비슷한 길이의 문서들을 인코더에게 전달한다. 문서 representation이 생성되면 디스크에 각 차원마다 32bit나 16비트로 저장하고 랭크를 계산하거나 인덱싱할때 불러와 사용한다.

Top-k Re-ranking

ColBERT는 reranking과 end-to-end retrieval에 모두 사용할 수 있다. Reranking은 주어진 쿼리 q에 대해 문서 셋이 작다. 문서 셋이 작기 때문에 메모리에 색인된 문서 representation을 올려 각 문서를 matrix of embedding으로 표현한다.

주어진 쿼리로 Eq를 계산함과 동시에 문서 셋 matrix로 이루어진  3차원 tensor D를 만든다. 배치 작업에서 문서들을 최대 문서 길이로 패딩했기 때문에 이 tensor \(D\)를 GPU로 옮기고 \(E_q\)와 \(D\)에 대해 배치 dot-product를 계산한다. 결과로 나온 3차원 텐서는 \(q\)와 각 문서에 대한 cross match matrices이다. 점수 계산은 문서 term에 대해서 max-pool로, 쿼리 term에 대해서는 합쳐서 reduce한다. 이렇게 계산한 점수로 정렬해 k개의 문서를 뽑는다.

ColBERT는 다른 neural ranker(특히 BERT 기반)들과 다르게 전체 비용이 계산된 임베딩을 모으고 옮기는 것에 지배된다. 일반적인 BERT ranker는 \(k\)개의 문서 랭킹시 쿼리 \(q\)와 문서 \(d_i\)로 이루어진 \(k\)개의 input을 주어야 하기 때문에 비용이 quadratic하게 늘어난다. 하지만 ColBERT는 하나의 쿼리 input만 주기 때문에 훨씬 싸다.

End-to-End Top-k Retrieval with ColBERT
End-to-end Retrieval의 경우 Top k개 대비 문서 셋 N이 상당히 크다(1000만개 이상). Late interaction 적용시 \(\text{MaxSim}\)에서 pruning을 수행한다. 쿼리 임베딩과 모든 문서에 대해서 MaxSim을 적용하지 않고 vector similarity를 이용해 쿼리와 관련 있는 문서들만 찾아 계산한다. 이를 위해 faiss와 같은 라이브러리를 사용한다. 그래서 faiss에 각 문서 임베딩과 문서 매핑 및 문서 임베딩 인덱스를 넣어야 한다.

먼저 faiss의 인덱스에서 vector-similarity 검색을 Nq번 수행해(Eq 내 임베딩 각각에 대해 수행) 각각 문서 top \(k^{'}\)개를 뽑아낸다(\(k^{'} = \frac{k}{2}\)). 그러면 총 \(N_q \times k^{'}\) 의 문서가 후보에 오르게 되고(겹치는 것도 있겠지만) 쿼리 임베딩과 높은 유사도를 가진다. 이 문서들을 대상으로 앞서 다뤘던 reranking을 수행한다.

Experiment

• Reranking의 경우 MRR은 BERT에 비해 조금 낮지만 latency는 훨씬 좋다.
• End-to-end retrieval은 BM25보다 latency가 많이 느리지만 높은 MRR을 보여주고 있다.

Reference

[1] Blog: [[논문 리뷰]ColBERT, ColBERTv2](https://pangyoalto.com/colbertv1-2-review/)
[2] Github: [ColBERT](https://github.com/stanford-futuredata/ColBERT)
[3] [Stanford C224U](https://web.stanford.edu/class/cs224u/slides/cs224u-neuralir-2023-handout.pdf?ref=pangyoalto.com)

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