Approximate nearest neighbor Negative Contrastive Learning for Dense Text Retrieval
1. Introduction
구글, 네이버 등에 검색을 하면 관련된 블로그, 카페, 위키 페이지가 검색 결과로 나옵니다. 이런 것들을 가능케 하는 것이 retrieval입니다. Retrieval는 텍스트가 입력되었을 때 해당 텍스트와 관련 있는 내용의 문서들을 찾는 모델입니다.
이런 retrieval에는 보통 겹치는 단어를 기반으로 탐색하는 sparse retrieval가 많이 사용되었습니다. 이는 입력 텍스트에 있는 단어가 가장 많이 쓰인 문서를 관련 문서로 선택하는 방법입니다. 그러나 이 sparse retrieval는 겹치는 단어를 보는 방식 때문에 겹치는 단어가 존재하지 않는 관련 문서는 찾아올 수 없는 태생적인 한계를 갖고 있습니다.
그래서 연구된 방식이 딥러닝을 활용한 dense retrieval입니다. 딥러닝 모델을 이용해 텍스트와 문서들의 정보를 추출한 뒤, 이들을 비교하여 가장 연관성 있는 문서들을 선택하는 것입니다. 이 방식은 텍스트들을 학습을 통해 변형 시킬 수 있기 때문에 겹치는 단어가 존재하지 않는 문서도 찾아올 수 있다는 장점을 갖습니다.
그러나 dense retrieval는 sparse retrieval보다 효용성이 떨어지는 단점을 갖고 있습니다. BERT의 등장으로 dense retrieval의 성능이 sparse retrieval의 성능을 뛰어넘기는 했지만, dense retrieval는 sparse retrieval에 비해 드는 자원이 너무 많고, 그에 비해 성능 향상 폭도 그렇게 크지 않기 때문입니다.
Dense retrieval의 학습 방식은 일반적인 분류나 회귀 문제와는 조금 다릅니다. Dense retrieval는 수백만 개의 문서들 중 입력된 텍스트와 가장 관련 있는 문서를 찾아내야 합니다. 이를 위해서 Dense retrieval는 관련 문서(positive)와 관련 없는 문서(negative)를 구분하는 식으로 학습합니다.
Positive document는 기준이 명확합니다. 입력 텍스트(질문)에 대한 정답 단어를 포함한 문서를 찾으면 됩기 때문입니다. 그러나 negative document는 그 외의 모든 문서가 됩니다. Negative 문서 후보가 굉장히 많기 때문에 적절한 negative 문서를 선정하는 것은 매우 어려운 일입니다. 그렇기 때문에 Dense Retrieval 학습에 가장 중요한 포인트는 적절한 negative 문서를 어떻게 선택하냐입니다.
Dense retrieval에서 negative 문서를 샘플링하는 방법을 연구했던 논문으로 DPR(Dense Passage Retrieval for Open-Domain QA)이 있습니다. DPR에서는 in-batch negative를 사용해 학습을 수행합니다. 딥러닝 학습은 보통 여러 개의 데이터를 묶어서 배치 단위로 진행합니다. In-batch negative는 이 하나의 배치 안에 있는 다른 질문의 관련 문서를 negative로 취급하는 방식입니다. 또, 성능을 더 높이기 위해 sparse retrieval의 한 종류인 BM25를 이용해 단어가 겹치지만 관련 문서가 아닌 'hard negative'를 추가합니다.
그러나 이 방식들은 한계가 존재합니다. 아래 그래프는 입력 텍스트(query)와 문서들을 임베딩 공간에 매핑한 모습을 그린 그래프입니다.
그래프를 보면 DPR의 negative sampling 방식이 최선이 아니란 것을 알 수 있습니다. 직관적으로 생각을 해봤을 때, dense retrieval의 성능을 높이기 위해선 모델이 가장 positive 문서로 착각을 하기 쉬운 문서를 negative로 선택하는 것이 필요합니다. 그러나 위 그래프를 보면 DPR의 BM25를 이용한 negative(BM25 Neg)나 in-batch negative(Rand Neg, 배치는 랜덤으로 결정되기 때문에 사실상 랜덤과 같습니다.)와 query 사이의 거리가 생각보다 먼 것을 확인할 수 있습니다.
ANCE는 negative로 이들보다 모델이 실제로 헷갈려 하는, 실제 dense representation 공간에서 positive 문서와 가까운 거리에 있는 문서들을 negative로 설정해 학습하는 방식입니다. 위 그래프로 보자면 DR Neg들을 사용해 모델을 학습시키는 것입니다. 이를 통해 DR 모델을 더 효율적으로 학습 시키자 하는 것이 ANCE의 핵심 내용입니다.
2. 문제 정의
논문에서 전개하는 수식을 이해하는데 사전지식이 굉장히 많이 필요한 것 같습니다. 이들을 완전히 해석하기 보단 제가 직관적으로 느낀 점들을 토대로 설명을 해보겠습니다. 어쨌거나 여기서 설명하고자 하는 것은 더 좋은 negative 문서를 선택하는 것이 모델 학습에 도움이 된다는 것을 증명하는 것입니다.
관련 문서를 찾고자 입력하는 텍스트를 query $q$라고 하고, 문서 집합을 corpus $C$라고 합니다.
이 때 dense retrieval의 목적은 $q$와 관련된 positive 문서 $D^+=\{d_1,...d_n\}$을 $C$에서 찾아야 합니다.
$q$와 $d$는 BERT를 이용해 $g(q;\theta)$와 $g(d;\theta)$로 각각 인코딩 합니다.
관련 문서를 찾는 방법은 $q$와 $d$ 벡터의 유사도가 높을수록 관련 문서라고 판별합니다. 이 유사도를 측정하는 similarity function으로는 dot product나 cosine similarity가 주로 사용됩니다.
이를 학습하기 위해 positive 문서 $d^+$와 negative 문서 $d^-$를 활용한 binary cross entropy loss, hinge loss, negative log likelihood loss가 주로 사용됩니다.
$$ l(f(q,d^+),f(q,d^-)) $$
학습 목표는 이를 이용한 최적의 모델 파라미터 $\theta^*$를 계산해내는 것입니다. $\theta^*$은 이 loss function을 최소화하는 최적의 모델 파라미터 $\theta$를 말합니다.
$$ \theta^*=argmin_\theta \Sigma_q \Sigma_{d\in D^+} \Sigma_{d^-\in D^-} l(f(q,d^+),f(q,d^-)) $$
$\theta^*$의 계산 과정은 모델의 stochastic gradient descent(SGD)를 통해 이뤄집니다. DR 모델의 $t$번째 SGD step은 아래와 같이 쓸 수 있습니다.
$$ \theta_{t+1}=\theta_t-n{1\over N_{P_{d^-}}}\nabla_{\theta_t}l(d^+,d^-) $$
SGD step은 loss function을 모델의 파라미터 $\theta$로 미분한 값을 이용해 수행합니다. 여기서 $1\over{N_{P_{d^-}}}$은 gradient 값을 조절해주는 역할을 합니다.
이 SGD step 수식의 convergence rate(학습 수렴 속도)는 아래와 같이 식을 풀어 계산할 수 있다고 합니다.
여기서 마지막 (9)번 식을 보면 모든 항에 $E_{P_{D^-}}$가 포함되어 있는 것을 확인할 수 있습니다. 즉, "어떤 negative 문서를 사용하는가"가 convergence rate(학습 속도)에 영향을 미친다는 것입니다.
직관적으로 생각했을 때, 더 큰 gradient 수치를 발생시키는 negative는 모델이 우리가 원하는 optimal한(수렴하는) 지점에 더 빠르게 도달할 수 있도록 해줄 것입니다. 즉, 위의 결과와 함께 생각하면 gradient 수치를 크게 발생시키는 negative는 더 빠른 학습 속도를 낼 수 있도록 해준다는 것입니다.
그렇다면 이 gradient 수치는 언제 커질까요? gradient의 계산에는 loss의 미분값이 사용됩니다. 즉, loss가 커질수록 gradient도 커진다는 것입니다. 결과적으로 loss가 커지는 negative를 선정해야 한다는 것입니다. 그렇다면 DR 모델의 학습과정에서 loss가 커지는 negative는 어떤 문서일까요?
다시 위 그래프를 봅시다. 위 그래프는 DR이 document들을 임베딩 공간에 매핑한 것을 보여줍니다. 여기서 만약 노란색 BM25 Neg를 Query의 negative로 사용한다고 생각해 봅시다. query와 BM25 Neg 사이의 거리는 이미 꽤 멉니다. 즉, 모델이 이미 BM25 Neg는 negative로 어느정도 간주하고 있다는 것이죠. 이는 원래 모델이 생각한 negative와 학습으로 제공된 새로운 negative 사이의 거리가 작다라고도 볼 수 있습니다. 즉, loss가 작아지게 됩니다.
이는 In-batch negative 방식도 마찬가지입니다. in-batch negative는 사실상 랜덤한 document를 negative로 사용해 학습하는 것인데, corpus의 크기가 매우 크기 때문에 여기서 query와 거리가 가까운 document를 선택할 확률은 0에 가깝습니다. 즉, in-batch negative 방식도 큰 loss를 갖기는 힘듭니다.
여기까지 왔으면 저자가 하고 싶은 말이 무엇인지를 알 수 있습니다. 학습을 빠르게 시킬 수 있는, 큰 loss를 갖게 하는 negative는 DR 모델이 가장 positive와 헷갈리는 document라는 것이고, 이들을 negative로 사용해야 DR 모델이 더 빠르게 학습될 수 있다는 것입니다.
3. ANCE
그렇다면 관건은 DR 모델이 가장 positive와 헷갈려하는 negative, 즉, query와 거리가 가까운 negative들을 어떻게 찾아서 사용하느냐가 됩니다. DR 모델은 query와 corpus 내의 document들을 각각 임베딩 벡터로 변환한 뒤, 그 벡터들 간의 거리를 similarity function을 사용해 측정합니다. 즉, query와 가장 가까운 negative document들을 찾기 위해선 모든 벡터들을 계산해 두어야 한다는 것입니다.
그러나 corpus의 크기는 매우 크기 때문에 이들을 모두 계산하는데 시간이 오래 걸립니다. 게다가 게다가 모델은 매 학습 step마다 파라미터가 업데이트 되기 때문에 이들의 벡터도 매 step마다 달라지게 됩니다. 사실상 매 step마다 이를 계산하는 것은 불가능한 일이 됩니다.
그렇기 때문에 ANCE는 학습과 document들의 벡터 갱신을 동시에 진행합니다. 이 document들의 벡터가 모두 계산되는 k step마다 document의 벡터를 갱신하는 방법을 사용해 학습이 끊기지 않고 정상적인 속도로 진행될 수 있도록 하는 것입니다.
위 그림과 같이 k-1 시점에서 계산된 document 벡터를 이용해 k+1 지점까지 학습을 수행함과 동시에, k 시점의 모델 파라미터를 가지고 새로운 document의 벡터들을 계산을 수행합니다. k+1 시점이 되면, k 시점에서 계산된 document 벡터들로 갱신하여 다시 학습을 이어가는 방식입니다. ANN(Approximate Nearest Neighbor) index는 Faiss라는 효율적인 임베딩 벡터 인덱스 계산 라이브러리를 사용해 수행했다고 합니다. 또 학습과 document 벡터 계산에 GPU를 1:1 비율로 할당하여 수행했다고 합니다.
4. 실험
Dataset
- web search : TREC 2019
- OpenQA : Natural Questions, TriviaQA
- commercial search engine
Baseline
- 하나의 BERT 모델을 이용해 query와 document를 각각 임베딩하는 siamese 구조 사용.
- Negative 선정 방식으로 random sampling, BM25 sampling, BM25+random 방식들과 비교 실험.
- 모델의 성능을 높이기 위해 학습 초기에는 BM25를 활용해 negative를 샘플링해 학습하는 BM25 warm up 방식을 사용.
- ANCE negative는 ANN index의 top200의 document 중 랜덤으로 하나를 선정하여 사용.
결과
결과는 ANCE 방식이 성능 면에서도 기존의 DR 모델들이나 sparse retrieval보다 좋게 나타났습니다. 하지만 아직 BERT를 이용한 rerank 방식보다는 우월한 성능을 보이지 못했습니다. (BERT Reranker : 하나의 BERT에 query와 document를 함께 입력하여 attention을 통해 둘 사이의 관계를 파악하는 방식. 성능이 좋지만 similarity function을 사용하지 않기 때문에 벡터를 미리 계산해 둘 수 없어 속도가 매우 느리다.)
그렇지만 다른 DR 모델들에 비해서 Rerank 방식과 Retrieval 방식 사이의 점수 차이가 가장 적은 모습을 보입니다. 여기서 Rerank 방식은 BM25와 같은 sparse retrieval의 결과를 DR 모델을 이용해 다시 rank를 메기는 방식을 말하고, Retrieval는 처음부터 DR 모델을 이용해 관련 문서 rank를 메기는 방식을 말합니다. 기존 방식의 DR 모델들은 Retrieval 방식이 Rerank 방식보다 점수가 많이 떨어져 Rerank 방식을 사용하는게 좋았지만, ANCE는 sparse retrieval에 의존하지 않고도 좋은 성능을 낼 수 있다는 것을 보여줍니다.
다음으로 위 그래프를 한번 보겠습니다. 위 그래프는 DR 모델이 찾은 query와의 top-200 document들의 similarity score를 그린 것입니다. 보면 점수가 대체로 상위 50위 내의 문서들에 DR score가 치중되어 있다는 것을 확인할 수 있습니다. 즉, DR 모델의 핵심은 이 상위의 문서들을 정확하게 구분해낼 수 있도록 하는게 핵심이란 것을 다시 한 번 확인할 수 있습니다.
또 각 negative sampling 방식으로 test셋에서 negative를 선정했을 때 그 중 얼마나 많은 비율이 top-200에 속하는지를 확인해 보았습니다. Random sampling 방식은 top-200에 속하는 문서가 0%였고, BM25 negative 중에는 15%가 속해 있었다고 합니다. ANCE의 경우 63%에서 시작해서 학습이 진행됨에 따라 100%에 가까운 비율을 갖게 되었습니다.
마지막으로, ANCE 방식이 더 높은 loss와 gradient 수치를 갖는다는 것을 증명하기 위해 학습이 진행되는 동안의 gradient 수치 변화와 loss 수치 변화를 그래프로 그려보았습니다. 그래프를 보면 ANCE(파란 점선)의 loss와 gradient 수치가 다른 방식들에 비해 훨씬 크다는 것을 알 수 있습니다. 그리고 이는 앞에서 증명했듯이 DR모델로 하여금 더 효율적이고 빠른 학습이 가능하도록 해줍니다.
5. 결론
이렇게 ANCE에 대해서 알아보았습니다. 수식적으로 설명하는 부분을 완전히 이해할 수는 없었지만, 어쨌거나 핵심적으로는 적절한 negative sample을 선정하는 것이 Dense Retrieval 모델의 학습에 중요하다라는 것을 알 수 있었습니다. ANCE는 그 방식으로 DR 모델이 실제로 헷갈려 하는, query와의 DR score가 가장 높은 document들을 negative로 사용하는 방식을 사용했고, 이로 인해 기존의 in-batch negative나 BM25 negative보다 더 효율적이고 좋은 학습을 이끌어냈습니다.