Search
🥎

[논문리뷰] Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval

Created
3/7/2021, 12:26:24 PM
Tags
QA
논문 리뷰: ANN search를 이용해서 어려운 negative sample 을 비동기적으로 뽑아 학습해 보자

요약

이번 포스트에서 우리가 살펴볼 본 논문에서는 Dense Retrieval이 본 성능을 내지 못하는 주요한 원인이 학습과 추론간의 데이터 분포가 다르다는 것임을 주장합니다. 이 문제를 해결하기 위해서 본 논문에서는 학습시 사용되는 negative sample을 뽑는 새로운 방법(ANCE)을 제안하였습니다.
실험을 통해서 BERT-Siamese 기반의 DR 모델 구조에서 ANCE 방법이 다른 sparse, dense retrieval 기반의 baseline 모델 보다 우수함을 확인할 수 있었습니다. 또한 제안된 방법이 sparse-retrieval과 BERT-reranking을 순차적으로 적용한 방식 보다 100배의 속도 향상을 얻으면서도 거의 같은 수준의 성능을 보였습니다.
지금부터 저자들의 주장은 무엇인지, 대안으로 제시된 방법은 어떤 것인지 한번 살펴 보도록 하겠습니다.

Dense Retrieval (DR)

구체적인 설명에 앞서, DR이 무엇인지 간략하게 짚고 넘어가도록 하겠습니다.
실시간 서비스에서 어떤 질의가 들어왔을 때, QA테스크에 fine-tuning 된 BERT에 모든 위키피디아 문서를 입력으로 넣어 주어 하나하나 자세하게 살펴보는 방식은 불가능에 가깝습니다. 떄문에 보다 빠른 로직으로 주어진 어떤 질의에 관련있는 문서들을 추리고, BERT와 같은 정교하지만 무거운 연산들을 추려진 문서들에 적용하여 속도와 정확성 두마리 토끼를 모두 잡는 방법을 사용합니다. 이러한 방식을 Cascade Information Retrieval (IR) 또는 Two-stage IR 이라고 합니다. 하지만 이 방식의 문제는 IR의 성능이 전체 모델 성능의 상한선을 결정한다는 것 입니다. 아무리 정교한 모델의 성능이 높다고 해도, 관련성 있는 문서를 추리는 퀄리티가 떨어진다면 전체 성능은 결코 좋지 않을 것 입니다.
기존에 대표적으로 사용되던 BM25와 같은 sparse retrieval(SR) 방식은 빠르지만, lexical 한 정보만을 고려하기 떄문에 좋은 성능을 보이지는 못하였습니다. 이 한계점을 극복하기 위해 최근에는 document 와 query 의 representation 을 비교하여 관련 있는 대상들을 찾는 dense retrieval(DR)이 활발하게 연구되고 있습니다. 가장 대표적으로 많이 사용되는 방식은 Sentence-BERT에서 사용된 Siamese Network 입니다. Query와 Document 를 encoding 하는 두개의 Encoder(BERT) 가 있고, encoding 된 두 representation의 내적값을 similarity function으로 사용하였습니다. 학습은 실제 positive sample과 negative sample 들이 포함되어 있는 batch 내에서 positive sample이 가장 큰 값을 갖도록 cross-entropy를 이용해 학습을 시키게 됩니다.
실제 추론시에는 미리 문서의 representation 을 encoding 시켜 놓은 후에 Approximate Nearest Neighbor(ANN) index를 build 합니다. 이후 입력으로 들어온 query 만 encoding 하여 ANN index 에서 query의 representation 과 유사한 문서를 빠르게 찾는 방식을 사용합니다. faiss와 같은 툴들을 이용해서 주어진 representation 과 유사한 후보들을 찾는 것은 매우 빠르게 수행할 수 있습니다. 이와 같은 방법은 속도에 있어서 굉장히 큰 이점을 갖지만, 아직까지 DR은 SR을 같이 병행하여 사용해야 하는 등, 기대하는 만큼의 충분한 성능을 보여주지 못하였습니다.

왜 Dense Retrieval은 충분한 성능을 보여주지 못하는가

저자는 이 문제의 원인으로 “DR representation 을 배우기 위해 사용했던 negative sample 과 실제 추론 과정에서 비교하게 되는 대상의 분포 간의 심한 차이”을 지목하였습니다. 저자는 실제로 두 분포의 차이가 나는지 확인하기 위해 negative sample 들의 representation 을 t-SNE으로 시각화하여 분석을 수행하였습니다.
실험 결과 예상대로, 추론 과정에서 DR 모델이 처리하는 negative 들은 실제 관련이 있는 문서들과 꽤 가까이 위치해 있었습니다. 하지만 BM25 또는 random sampling 을 이용해 뽑혀진 negative 들은 실제 관련 문서들의 위치와는 크게 떨어져 있었습니다. 만약 이런 negative 들로 모델을 학습 한다면, 실제 관련이 있는 문서들과 잘못 retrieval 된 negative 간의 representation 분리는 절대로 이루어지지 않을 것입니다. 즉 IR의 분별력을 높이기 위해서는 단순히 어느정도 관련 있는지 없는지에 대한 정보 뿐만 아니라, 어느정도 관련이 있는 것들 중에서 진짜 관련이 있는 문서를 고르는 능력 역시 배워야 한다는 것을 보여줍니다.

Approximate nearest neighbor Negative Contrastive Estimation(ANCE)

직관적으로 생각해 보면 모델이 어려워하는 negative sample 들을 이용해서 DR을 학습시키게 되면 모델은 더욱 분별력을 갖게 됩니다. 하지만 기존 방법들처럼 큰 코퍼스에서 random sampling을 통해 뽑은 대부분의 sample 들은 query 와 관련이 없는 문서들이기 때문에, 어려운 샘플을 뽑을 가능성이 거의 없습니다. 또한 본 논문의 목표는 sparse retrieval 을 넘어서는 모델을 만드는데 negative sample 을 만드는데 SR을 사용하는 것은 적합하지 않습니다.
위와 같은 문제를 해결하기 위해서 본 논문에서는 negative sample 을 만드는 Approximate nearest neighbor Negative Contrastive Estimation(ANCE)이라는 새로운 방법을 제안합니다. 이 방법은 학습된 모델로 만든 ANN index 로부터 retrieval 된 결과를 이용해서 어려운 negative sample을 만드는 방법입니다. ANN을 사용해 negative 를 만드는 과정은 실제 추론시에 학습된 모델이 사용되어지는 방식과 동일하기 때문에 기존의 분포가 일치하지 않던 문제를 해결할 수 있습니다.
저자들은 이 방법을 통해서 학습과 추론간의 데이터 분포 불일치(discrepancy)를 해결할 수 있다고 주장하였습니다. 즉, 모델이 실제 정답에 근접한 negative sample 들과 실제 정답간의 구분을 할 수 있도록 함으로서, 보다 어려운 문제를 푸는 과정을 통해 모델의 성능을 향상시킬 수 있는 방법인 것이죠.

비동기 학습 방법 : Asynchronous Training

Negative sample 뽑은 때 사용하는 ANN index 를 만들기 위해서는 다음과 같은 두가지 스텝이 필요로 합니다.
1.
Inference: 업데이트된 새로운 encoder 를 이용해서 전체 문서의 representation 을 계산합니다.
2.
Index: 전체 representation 에 대해서 ANN index 를 build 합니다.
하지만 ANN index 를 한번 업데이트 하기 위해서는 모든 문서를 BERT로 encoding해야 하기 때문에, 현실적으로 학습의 매 step(batch)마다 위 두 과정을 수행하는 것은 불가능합니다. 때문에 본 논문에서는 ANN index 를 update 하는 과정과 학습 과정을 병렬적으로 수행하는 비동기 학습 방법을 제안합니다. 방법은 다음과 같습니다.
1.
encoder 의 마지막 체크포인트 를 이용해 전체 corpus에 대해서 representation 을 계산하는 inference 를 수행하고 동시에 병렬적으로 trainer 는 게속  을 이용해서 학습을 수행합니다.
fk
ANNfk−1
2.
를 이용한 inference 가 끝나면 를 다시 indexing 하고, trainer 에서 학습중인 negative sample 들을 새로운  을 이용해서 뽑아 학습 코퍼스를 업데이트 합니다.
fk
ANNfk
ANNfk
즉 한쪽에서는 계속 ANN을 업데이트를 하고, 한쪽에서는 실제 모델을 업데이트 하다가, ANN 업데이트가 다 되면 negative sample 들을 다시 갱신하여 모델이 점점 어려운 문제를 풀게 되는 시스템으로 이해할 수 있겠습니다.

Experiment

평가는 TREC 2019 Deep Learning Track 데이터셋을 사용해서 진행하였으며, 이 데이터셋은 Bing 으로 들어온 백만개 이상의 질의에 대해서 관련된 passage 가 annotation 되어 있는 데이터셋 입니다. 이 데이터셋을 선택한 이유로 크고, 최신이고, 가장 현실적인 상황을 잘 반영했기 떄문에 사용했다고 저자는 기술하였습니다. 평가 메트릭은 MRR과 Recall@1k, NDCG를 사용하였습니다. 대부분의 성능은 retrieval 에 대한 성능을 측정하였고, 추가적으로 주어진 100개의 candidate 내에서 DR 모델을 이용해 rerank 하는 능력도 같이 검증하였습니다. (표에서 Rerank라고 나와 있는 부분)
DPR과 동일하게, 도메인의 제한이 없는 QA데이터셋인 OpenQA task 데이터셋으로도 평가를 수행하였습니다. 평가 방식은 Top-N안에 실제로 우리가 타겟팅 하는 passage 가 포함되어 있는지 아닌지 평가하는 메트릭을 사용하였습니다.
Table1 표 좌측을 보면, 학습에 사용한 negative sample 들의 variation 을 확인할 수 있습니다. 그리고 최종적으로 본 논문에서 제안한 방식으로 negative sampling 을 했을 떄의 성능을 가장 아래에 기록하였습니다. 결과적으로 sparse bag-of-words retrieval 없이도 ACNE 가 BM25 기반의 학습보다 상대적으로 20% 높은 성능 향상을 보여 주었습니다.
DR에서의 학습 방법의 variation 중에서는 BM25+Random Negative 가 sparse retrieval 의 성능보다 높은 성능을 기록하였는데, 여러 연구들에서 이렇게 학습된 DR 모델이 단순한 sparse retrieval 모델보다 좋은 성능을 갖지 못함을 언금하였습니다. 때문에 저자는 같은 BERT-Siamese 구조의 모델 중에서 ANCE만이 sparse retrieval 기법보다 robust 하고 우수한 성능을 갖는 모델은 ANCE 가 유일하다고 주장하였습니다.
주어진 N개의 document 를 rerank 하는 테스크에서도 제안하는 ANCE가 기존 DR또는 SR보다 우수한 성능을 보여 주었으며, 특히 TREC DL Document rerank 테스크에서는 BERT 를 사용한 방법보다 높은 성능을 기록함을 확인할 수 있었습니다. 이처럼 제안하는 ANCE 방법은 대부분의 평가 metric 에서 기존 baseline 모델의 성능을 향상시킬 수 있음을 보여주며, 제안된 모델의 우수함을 주장하였습니다.

ANCE Retrieval의 학습 및 추론 효율성

저자는 제안하는 ANCE 모델의 학습 또는 inference 시의 효율성에 대해서 분석을 진행하였습니다. Table 3에서는 실제 inference 시에 발생하는 시간을 측정하여 기록하였는데, 최종적으로 inference 시에는 쿼리 하나당 11.6ms 밖에 걸리지 않음을 보여주었고, BERT rerank를 기반으로 하는 Cascade Retrieval 방식(1.42s) 보다 약 100배 빠른 성능임을 시사하고 있습니다. 즉 성능도 뛰어난데, 속도 마저 100배 빠르다는 것을 강조하였습니다.

세부 사항

학습시 배치 내의 Positive, negative 비율은 1:1 이며, 각 query 에 대해 retrieve 된 document top 200 개중 한개를 random 하게 선택하여 negaitve sample로 사용하였습니다. 학습에 사용한 GPU는 총 8개이며 ANN encoding 과 학습에 배정된 GPU 갯수도 1:1 비율로 배정하였습니다. ANN index update는 10k step 마다 실시하며, 학습에 사용하는 batch-size 는 8입니다. 특이한 점은 step 2 마다 gradient accumulation 을 진행하였습니다. lr 는 5e-6을 사용하였으며 학습시 Optimizer 로 LAMB를 사용하였습니다.
저자는 ANN index 를 update 하는 시점에 대해서도 실험을 실시하였는데, 5K마다 update 를 하는 경우에 가장 smooth 하게 수렴하지만, inference 에 2배의 GPU가 필요함을 지적하였습니다. 때문에 같은 computation 을 사용했을 때에는 10k를 사용함이 가장 우수한 성능을 보여주었다고 주장합니다.

결론

ANCE fundamentally eliminates the discrepancy between the representation learning of texts and their usages in dense retrieval. Our ANCE trained dense retrieval model, the vanilla BERT-Siamese, convincingly outperforms all dense retrieval and sparse retrieval baselines in our large scale document retrieval and passage retrieval experiments. It nearly matches the ranking accuracy of the state-of-the art cascade sparse retrieval and BERT reranking pipeline. More importantly, all these advantages are achieved with a standard transformer encoder at a 1% online inference latency, using a simple dot-product in the ANCE-learned representation space.