Pre-training Tasks for Embedding-based Large-scale Retrieval

[접기]

목차

  • 1. Information Retrieval
    • 1.1. Scoring Function $ f $
    • 1.2. IR의 2단계
  • 2. 기존 모델들
    • 2.1. BM25
    • 2.2. Transformer 기반 모델
  • 3. 논문에서 제안하는 모델
  • 4. Pre-training Task


Pre-training Tasks for Embedding-based Large-scale Retrieval
- ICLR 2020. 04. -
https://arxiv.org/pdf/2002.03932.pdf

1. Information Retrieval


1.1. Scoring Function $ f $


일반적으로 IR(Information Retrieval) 문제에 사용되는 모델들은 다음과 같은 함수를 학습한다고 볼 수 있습니다.
$$ f : Q \times D \to \mathbb{R} $$여기서 $ Q $는 질의(query)들의 집합, $ D $는 문서(document)들의 집합입니다. 즉, 모델은 $ (q, d) \in Q \times D $에 대해 $ f(q, d) $라는 점수를 생성합니다. 이 점수는 $ (q, d) $가 서로 관련있는 질의와 문서 쌍인 경우 높아야하고, 관련 없는 쌍인 경우 낮아야합니다. 물론, 여기서 질의와 문서라고 부르기는 했지만 실제로는 다양한 정보 검색 문제들을 이 형태로 표현할 수 있습니다. 상품 추천 시스템에서는 $ q $가 사용자의 정보(최근에 산 물건, 찜목록 등)이고 $ d $가 후보 상품이 될 수 있습니다. 문서마다 적절한 태그를 붙이고자하는 검색엔진 개발자한테는 $ q $가 실제 웹 문서이고 $ d $가 후보 태그가 될 수 있습니다. 뒤에서 조금 더 자세히 살펴보겠지만, open-domain QA에서는 $ q $가 주어진 질의, $ d $가 정답을 포함하고 있는 단락을 나타낼 수 있습니다.

1.2. IR의 2단계


그러나 좋은 성능을 위해선 $ f $가 복잡하고 느릴 것이고, 이는 수십만 또는 그이상의 대규모 문서들을 대상으로 IR을 수행할 때 큰 문제가 됩니다. 즉, 실제 검색엔진이나 대규모 추천시스템에 적용할 수 없다는 뜻입니다. 그렇기 때문에 일반적으로 IR은 2단계를 거치는데, 첫 번째 단계에서는 빠르게 동작하지만 높은 Recall을 가지는 $ f $를 적용하여 후보 문서를 줄이고(top 1000이나 top 100 등) 두 번째 단계에서는 성능이 좋지만 느릴 수 있는 $ f $를 적용합니다. 그러나 최근에 이루어지고있는 대부분의 연구들은 두 번째 단계에서 성능을 향상시키려하고 있는데, 여기에서 성능이 아무리 올라도 첫 번째 단계에서 놓친 정답 문서들 찾아올 수는 없습니다. 이에 저자들은 첫 번째 단계에 집중하는 것이 중요함을 이야기하면서 자신들의 연구를 소개합니다.

2. 기존 모델들


2.1. BM25


앞서 언급한 $ f $의 적절한 사례 중 하가나 바로 BM25입니다. BM25는 2009년에 발표되었고[1] 아직까지도 이기기 힘든 강력한 알고리즘으로 인정받고 있습니다. IR의 2단계 방식에서 첫 번째 단계에 대부분의 연구들이 사용하는 알고리즘이기도 합니다. BM25는 결국 TF(Term Frequency)와 IDF(Inverse Document Frequency)를 바탕으로 사람이 설계한 함수이며, 때문에 원하는 task에 맞추어 최적화하기 힘들다는 단점이 있습니다.

2.2. Transformer 기반 모델


최근에는 BERT와 그 후속 연구들(XLNet, RoBERTa 등)이 Transformer를 기반으로 구성된 $ q $, $ d $간의 cross-attention 모델을 제안하여 높은 성능을 내고 있습니다. 특히 IR의 2단계 방식에서 두 번째 단계에 말입니다. 그러나 이는 질의 $ q $가 들어올 때마다 모든 문서 $ d $와의 cross-attention을 계산해야하기 때문에 후보 문서가 많을수록, 각각의 문서가 길수록 속도가 매우 느리다는 단점이 있습니다.

3. 논문에서 제안하는 모델


3.1. 구조


논문에서 제안하는 two-tower 모델(왼쪽)과 BERT 계열의 cross-attention 모델(오른쪽)
저자들이 제시하는 모델은 위 그림에서 왼쪽의 형태로, 질의와 문서가 서로 독립된 모델 $ \phi(\cdot) $와 $ \psi(\cdot) $를 각각 거친다는 점이 기존에 연구된 BERT 계열의 cross-attention 모델과 다른 점입니다. $ \phi(q) $와 $ \psi(d) $는 각각 질의와 문서의 고정 크기 벡터 표현(논문에서는 embedding이라고 부름)을 만들어냅니다. 이에 따라 scoring function $ f $는 다음과 같이 두 벡터의 내적으로 정의합니다.
$$ f(q, d) = \langle \phi(q), \psi(d) \rangle $$그리고 $ \phi(\cdot) $와 $ \psi(\cdot) $는 각각 독립된 transformer들을 쌓아서 구성합니다. BERT에 익숙한 분들은 쉽게 생각해서 BERT의 cross-attention을 포기하고 질의와 문서 각각이 독립적으로 self-attention을 계산한다고 보면 될 듯 싶습니다. 실제로 모델의 입력이 [CLS], tok1, ..., tokn [SEP] 형태로 구성된다는 점과 최종 score 생성을 위해 [CLS] 위치의 벡터를 사용한다는 점에서 모델 구조가 BERT와 거의 동일합니다.

4. Pre-training Task


그리고 또 하나 논문에서 강조하는 점은 어떤 pre-training task를 사용했는지에 따라 모델의 성능 차이가 크다는 점입니다. 실제로 two-tower 모델에서는 BERT에서 사용한 MLM(Masked Language Model)로는 큰 성능향상을 하지 못했지만 ICT(Inverse Cloze Task), BFS(Body First Selection), WLP(Wiki Link Prediction)을 사용했을 때에는 성능이 많이 향상된다는 점을 아래 표와 같이 결과로 보여줍니다.
train/test 데이터 비율에 따른 SQuAD 데이터셋에서의 [email protected] 결과 (단위: %)
예를 들어, 마지막줄을 보면 MLM으로는 [email protected]이 80점이었지만 제시한 3가지 task를 모두 사용하면 97점으로 매우 향상됨을 볼 수 있습니다. 그리고 학습 데이터와 평가 데이터의 비율에 따라서도 비교를 했는데, 학습에 매우 적은 데이터만 사용했을 때에도 높은 성능을 얻었습니다. 위 표는 SQuAD 데이터셋에 대한 결과이며, 논문을 보면 Natural Questions 데이터셋에 대해서도 동일한 실험을 하고 결과를 제시하고 있습니다.