1. 개요
ODQA(Open Domain Question Answering)와 기존의 해결방식
Open Domain Question Answering은 질문이 주어졌을 때, 해당 질문의 답을 포함하는 문서를 대량의 문서집합으로부터 찾아서, 해당 문서에서 질문에 알맞은 답을 추출해내는 task입니다. 이 논문이 출시되었을 당시엔 BERT를 활용한 방식이 SOTA를 달성한 상태였습니다. 그러나 BERT를 사용하는 방식은 기존 방식들에 비해 너무 느리다는 문제가 있었습니다.
BERT 방식이 기존 방식들보다 느린 이유는 질문과 관련 있는 문서를 찾는 방식에 있습니다. BERT는 질문과 문서 사이의 관련도(유사도)를 측정하기 위해 아래와 같은 과정을 거칩니다.
질문(Query)과 문서(Document)를 하나의 BERT에 같이 입력함으로써 attention 연산을 통해 질문과 문서의 글자들을 하나한 비교하는 방식을 사용하는 것입니다. 질문과 문서가 함께 입력되어야 하기 때문에 만약 백만개의 문서 중에서 질문과 관련 있는 문서를 찾고자 한다면 백만 번의 BERT 연산이 필요하게 되는 것입니다.
효율성을 높이기 위한 관건은 질문과 문서의 임베딩 벡터를 따로 계산할 수 있는가 입니다. 만약 질문과 문서를 BERT와 같이 한번에 입력받지 않고 따로 계산할 수 있다면, 백만개 내지 천만개가 되는 문서들의 임베딩 벡터를 미리 계산할 수 있게 됩니다. 미리 계산해 두기 때문에 연산 속도가 확연히 줄어들 수 있죠.
이런 방식을 채용한 것이 아래와 같은 방식입니다.
인코더를 사용해 질문과 문서를 따로 임베딩한 뒤에, 두 임베딩 벡터 사이의 유사도 점수를 계산하는 방식입니다. 여기서 유사도 계산에는 cosine similarity나 dot-product와 같은 연산이 사용됩니다. 그러나 이 방식은 앞서 본 BERT all-to-all interaction 방식에 비해 성능이 낮습니다.
ColBERT는 이 2가지 방식의 장점을 모두 수용하는 모델을 만들고자 진행된 연구입니다. ColBERT의 핵심은 "Late interaction"입니다. 그럼 지금부터 ColBERT에 대해서 알아보도록 하겠습니다.
2. 구조
2-1. 인코더
ColBERT는 기본적으로 질문과 문서를 따로 인코딩합니다. 여기에 BERT가 사용됩니다. 질문과 문서를 따로 인코딩하지만 같은 BERT가 사용됩니다. 즉, 하나의 BERT로 질문과 문서를 따로 임베딩합니다. 그렇지만 질문과 문서 텍스트는 구분되어야 하기 때문에 새로운 special 토큰 "[Q]"와 "[D]"를 사용합니다. 이 토큰은 BERT에 들어오는 텍스트가 질문인지 문서인지를 나타냅니다. 이 special 토큰들은 맨 앞의 [CLS] 토큰 바로 뒤에 위치합니다.
(ColBERT input 예시)
질문(query) : "[CLS] [Q] 세상에서 가장 큰 절은? [MASK] [MASK] [MASK]"
문서(document) : "[CLS] [D] 불국사는 경상북도 경주시 토함산에 있는 대한불교 조계종 소속 사찰이다. [PAD]"
또 질문을 인코딩할 때는, 문서와 다르게 조금 특이한 방식의 query augmentation이 적용됩니다. 보통 텍스트 전처리를 할 때 배치 단위의 연산을 위해 길이가 제각각인 문장들을 [PAD] 토큰을 이용해 길이를 맞춰주는 작업을 수행합니다. 그러나 여기선 [PAD] 토큰 대신 [MASK] 토큰을 사용합니다.
논문에 따르면, 이 [MASK] 토큰을 대신 사용함으로써 질문에 사용된 단어들의 정보를 확장하거나, 각 단어들의 중요성에 따라 가중치들을 재조정할 수 있게 된다고 합니다.
설명이 복잡한데, [PAD] 토큰을 사용했을 때와의 차이점을 생각해 보겠습니다. BERT 연산을 수행할 때, [PAD] 토큰의 값들은 모두 무시됩니다. 그러나 이를 [MASK]로 대체한다면 모델은 "'[MASK]'에 뭐가 들어가야하지?"를 고민하면서 이 부분에 더 의미있는 값을 부여하게 될 것입니다.(BERT는 그렇게 pre-training되었기 때문에) 그러니 [PAD]를 사용했을 때보다 '같은 문장이더라도 더 많은 의미를 내포한 representation을 추출할 수 있다' 라고 생각한 것 같습니다.
너무 복잡하면 그냥 이렇게 했는데 성능이 더 좋더라 정도로 받아들여도 좋을 것 같습니다. 어쨌든 결과가 중요하니깐요...
BERT의 output은 다시 linear layer를 통해 차원 수를 m으로 조절합니다. 차원 수를 줄임으로써 더 효율적인 유사도 계산이 가능하도록 하는 것입니다.
그 후, L2 norm을 사용해 값들의 범위를 조절합니다. 범위를 조절함으로써 두 임베딩 벡터의 내적을 계산했을 때 그 결과값의 범위가 -1에서 1 사이로 분포되도록 합니다.
문서의 경우 보통 텍스트 길이가 매우 깁니다. 그렇기 때문에 조금이라도 더 효율적인 계산을 위해 punctuation mark(마침표, 느낌표 등)에 해당하는 임베딩 값들은 문서 맥락에 크게 중요하지 않다고 판단하여, 모두 필터링을 한다고 합니다.
인코딩 과정을 수식으로 표현하면 아래와 같이 쓸 수 있습니다.
$ E_q:=Normalize(CNN(BERT("[Q]q_0q_1...q_l[mask]...[mask]"))) $
$ E_d:=Filter(Normalize(CNN(BERT("[D]d_0d_1...d_n")))) $
여기서 뜬금없이 CNN이 나타나는데, 사실 linear layer는 kernel 크기가 1인 convolution filter 연산과 같다고 볼 수 있습니다. 즉, CNN은 linear layer를 조금 더 일반화한 유형이라고 볼 수 있습니다. 여기선 그냥 linear layer로 해석해도 될 것 같습니다. 코드에서도 아래와 같이 linear 레이어를 사용했으니까요.
# punctuation mark의 임베딩값들은 필터링 합니다.
self.skiplist = {w: True
for symbol in string.punctuation
for w in [symbol, self.raw_tokenizer.encode(symbol, add_special_tokens=False)[0]]}
def mask(self, input_ids, skiplist):
mask = [[(x not in skiplist) and (x != self.pad_token) for x in d] for d in input_ids.cpu().tolist()]
return mask
def query(self, input_ids, attention_mask):
input_ids, attention_mask = input_ids.to(self.device), attention_mask.to(self.device)
Q = self.bert(input_ids, attention_mask=attention_mask)[0]
Q = self.linear(Q)
mask = torch.tensor(self.mask(input_ids, skiplist=[]), device=self.device).unsqueeze(2).float()
Q = Q * mask
return torch.nn.functional.normalize(Q, p=2, dim=2)
def doc(self, input_ids, attention_mask, keep_dims=True):
D = self.bert(input_ids, attention_mask=attention_mask)[0]
D = self.linear(D)
mask = torch.tensor(self.mask(input_ids, skiplist=self.skiplist), device=self.device).unsqueeze(2).float()
D = D * mask
D = torch.nn.functional.normalize(D, p=2, dim=2)
return D2-2. Late interaction
위와 같이 인코딩된 질문과 문서 벡터는 late interaction을 통해 두 벡터의 유사도를 계산하게 됩니다. Late interaction은 다음과 같이 계산됩니다.
1. squared L2나 cosine similarity를 통해 두 벡터의 유사도 matrix를 계산.
2. 유사도 matrix는 3-dimensional tensor가 되는데, 이를 document 축으로 max-pooling을 적용하고, query축으로 summation을 다시 적용하여 1-dimensional tensor 형태의 최종 점수를 출력.
$ S_{q,d}:=\Sigma_{i\in [|E_q|]}max_{j\in[|E_d|]}E_{q_i}\cdot E^T_{d_j} $
즉, 기존의 dot-product와 같은 단순 similarity 연산에서 한 발 더 나아가, max-pooling과 summation을 다시 적용한다는 것이 차별점이라고 볼 수 있습니다. 코드 상으로는 아래와 같이 구현할 수 있습니다.
def score(self, Q, D_padded, D_mask):
'''
Q: torch.Tensor(size=(Batch_size, seq_length, m))
D_padded: torch.Tensor(size=(Batch_size, seq_length, m))
'''
if self.colbert_config.similarity == 'l2':
l2_dist = -1.0 * ((Q.unsqueeze(2) - D_pasdded.unsqueeze(1))**2).sum(-1)
s = l2_dist.max(-1).values.sum(-1)
return s최종적으로 모델은 질문($q$), 관련 있는 문서($d^+$), 관련 없는 문서($d^-$)가 <$q, d^+, d^-$>로 주어졌을 때, query와 $d^+$, $d^-$ 사이의 score를 각각 계산한 뒤, 해당 score들의 cross-entropy를 통해 loss를 계산하게 됩니다. (q와 $d^+$ 사이의 score는 1, q와 $d^-$ 사이의 score는 0이 되도록)
2-3. Top-k Re-ranking with ColBERT
이 분야에서 retrieval를 활용하는 방법은 re-ranking 방식과 end-to-end 방식으로 나뉩니다. re-ranking 방식은 BM25와 같은 방식으로 top-k 문서를 추출한 뒤, 이들을 retrieval를 사용해 다시 순위를 메기는 방식입니다. 이 경우 이미 top-k개로 문서의 수가 크게 줄어든 상태이기 때문에 질문 query와 모든 document 사이의 유사도를 계산해도 시간이 크게 오래 걸리지 않습니다.
수행 방식은 다음과 같습니다.
1. query의 임베딩 벡터의 길이를 document 임베딩 벡터와 맞추기 위한 패딩을 수행한다.
2. query 임베딩 벡터와 top-k개의 document 임베딩 벡터 사이의 dot-product를 수행한다.
3. 결과에 max-pooling과 summation을 적용하여 점수를 출력.
4. 출력된 점수를 기반으로 top-k 문서들을 재정렬.
'''
top-k 문서들과 질문 사이의 유사도 계산하기
m : 임베딩 벡터의 차원 수
q_len : query의 길이
d_len : document의 길이
'''
Q = self.bert(Query) # shape=[1, q_len, m]
D = self.bert(Documents) # shape=[k, d_len, m]
# Document의 길이로 padding
Q = pad_to_document_length(Q, length=d_len) # shape=[1, d_len, m]
# dot-product
score = torch.matmul(Q.repeat(k, 1, 1), D.transpose(1, 2)) # shape = [k, d_len, d_len]
score = score.max(-1).values.sum(-1) # shape = [k]2-4. End-to-end Top-k Retrieval with ColBERT (with faiss)
End-to-end 방식의 경우, top-k가 주어지지 않고 주어진 문서 전체로부터 top-k개를 추려야 합니다. 그렇기 때문에 re-ranking 방식과 같이 모든 문서와의 유사도를 계산하기에는 너무 많은 시간을 필요로 합니다. 이를 효율적으로 수행하기 위해 facebook의 faiss를 사용합니다.
faiss는 유사 임베딩 벡터를 효율적으로 찾을 수 있도록 도와주는 일종의 알고리즘입니다. 이 방식은 전체 문서들의 임베딩 공간을 k-means clustering을 활용해 P개의 cell로 분할합니다. 각 cell에는 해당 cell의 평균값과 같은 벡터가 존재합니다. 전체 문서들은 이 cell의 평균값들과의 유사도를 계산하여 가장 유사한 cell들에 할당되게 됩니다.
여기서 top-k를 탐색하고자 할 때는 질문 query와 가장 유사한 p개의 cell에 대해서만 탐색을 진행합니다. 이를 통해서 탐색 시간을 크게 줄일 수 있게 됩니다.
3. 실험
실험에 사용된 데이터셋은 MS MACRO와 TREC 2가지입니다.
- MS MACRO Ranking : Microsoft가 만든 데이터셋으로 8백8십만개의 Web page text들을 passage로 하여, 백만개의 Bing의 질문 query들에 대해 대답해야 하는 task입니다. 각 질문들에 대해 정답을 담고 있는 gold passage가 하나씩 할당되어 있습니다. 모델 평가는 MRR@10을 통해 수행합니다.
- TREC CAR : Wikipedia의 2천9백만개의 passage로 이루어진 데이터셋입니다. 성능 테스트는 2,254개의 질문을 담고 있는 TRECT 2017 CAR로 수행했습니다.
learning rate는 $3*10^{-6}$으로 설정했으며, 배치 크기는 32로 설정했습니다. emgedding dimension ($m$)은 128로 설정했다고 합니다.
3-1. Top-k Re-ranking
MS MACRO의 결과는 위와 같이 나타났습니다. ColBERT는 BERT와 유사한 성능을 보이면서도, FLOP(소모 컴퓨팅 자원과 속도) 수치는 BERT를 사용하지 않는 가벼운 모델들과 거의 유사한 수치를 보입니다. 수많은 document들의 벡터들을 미리 계산할 수 있기 때문입니다. 또, BERT와 late interaction을 통해 BERT의 높은 성능도 확보할 수 있었습니다.
이런 모습은 TREC CAR에서도 똑같이 나타났습니다. ColBERT는 BERT를 활용한 것과 비슷한 점수(MAP)를 달성했습니다.
3-2. End-to-end Top-k retrieval
End-to-end는 기존의 BM25, doc2query 등과 같은 알고리즘 기반의 방식들과 비교를 했습니다. (ORQA, REALM 등의 연구와 비슷한 시기에 나와 이런 모델들과는 성능 비교가 되지 못한 것 같습니다.) Latency는 꽤 높게 나오지만 MRR@10 점수나 Recall(특히 top-k가 작을수록) 점수를 보면 ColBERT가 더 높은 점수를 달성한 것을 볼 수 있습니다.
3-3. Ablation studies
마지막으로 ColBERT에 사용된 방식들이 모델 성능에 어떤 영향을 주었는지에 대해서 살펴보고 마무리 해보겠습니다.
1) Late Interaction 방식의 효용성
위 그래프에서 [A] 모델은 단순히 dot-product를 사용해 질문과 문서 임베딩 벡터의 유사도를 계산한 모델입니다. Late interaction을 적용한 [D] 모델과 성능을 비교했을 때 점수가 확연히 차이나는 것을 확인할 수 있습니다. 더불어 max-pooling 대신 average-pooling을 적용한 모델 [B]도 점수가 크게 떨어지는 것을 볼 수 있습니다.
2) query augmentation ([mask] 토큰으로 패딩)
query augmentation을 사용하지 않은 모델은 [C]의 결과로, [D]의 결과와 점수 차이가 어느정도 나는 것을 확인할 수 있습니다.
3) End2end colBERT
re-ranking 방식의 ColBERT [E]와 end2end 방식의 ColBERT [F]을 비교했을 때 [F]의 점수가 더 높습니다. BM25과 같은 방식에 의존하지 않은 딥러닝 방식이 더 좋은 성능을 발휘한다는, DPR, ORQA, REALM과 같은 맥락의 분석입니다.
4. 결론
ColBERT는 BERT를 사용한 모델의 높은 성능과, BERT를 사용하지 않은 모델들의 효율성을 함께 가져가고자 하는 목표로 연구된 논문입니다.
SBERT와 같이 BERT를 이용해 각 query와 document를 따로 인코딩함으로써 BERT의 성능 상의 이점도 살리면서, 미리 계산해 둘 수 있는 offline-indexing을 통해 연산의 효율성도 살릴 수 있었습니다. 또, BERT의 성능을 따라잡기 위해 late interaction(max-similarity)을 적용하였고, query augmentation을 적용했습니다.
개인적으로 DPR의 in-batch negative 방식과 결합되어 사용하면 더 좋은 성능을 보일 것 같습니다. 나중에 실험해 보면 좋을 것 같습니다...
논문:
https://arxiv.org/pdf/2004.12832.pdf
github:
