[논문 리뷰] RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder

Shitao Xiao, Zheng Liu외 2명|arXiv (Cornell University)|2022. 05. 24.

Topic Modeling인용 수 7

한 줄 요약

RetroMAE는 검색 중심 언어 모델을 위한 사전 훈련을 위해 비대칭 설계를 사용하는 새로운 마스크된 오토인코딩 프레임워크를 제안한다. 이는 전체 인코더 BERT와 단일 레이어 디코더를 사용하며, 디코더에는 50–70%의 강력한 마스킹, 인코더에는 중간 수준의 마스킹(15–30%)을 적용한다. 이러한 설정은 인코더가 깊은 의미적 표현을 학습하도록 유도하여, 최소한의 데이터와 계산 자원으로 BEIR 및 MS MARCO 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Despite pre-training's progress in many important NLP tasks, it remains to explore effective pre-training strategies for dense retrieval. In this paper, we propose RetroMAE, a new retrieval oriented pre-training paradigm based on Masked Auto-Encoder (MAE). RetroMAE is highlighted by three critical designs. 1) A novel MAE workflow, where the input sentence is polluted for encoder and decoder with different masks. The sentence embedding is generated from the encoder's masked input; then, the original sentence is recovered based on the sentence embedding and the decoder's masked input via masked language modeling. 2) Asymmetric model structure, with a full-scale BERT like transformer as encoder, and a one-layer transformer as decoder. 3) Asymmetric masking ratios, with a moderate ratio for encoder: 15~30%, and an aggressive ratio for decoder: 50~70%. Our framework is simple to realize and empirically competitive: the pre-trained models dramatically improve the SOTA performances on a wide range of dense retrieval benchmarks, like BEIR and MS MARCO. The source code and pre-trained models are made publicly available at https://github.com/staoxiao/RetroMAE so as to inspire more interesting research.

연구 동기 및 목표

• 기존의 사전 훈련 방법이 밀도 있는 검색을 위한 강력한 문장 수준 표현을 개발하는 데에 한계가 있음을 해결하기 위해.
• 오토인코딩 기반의 검색 사전 훈련에서 데이터 효율성과 표현 품질 문제를 극복하기 위해.
• 입력 문장으로부터 훈련 신호를 최대한 활용하면서 동시에 인코더의 복원 난이도를 높이는 프레임워크를 설계하기 위해.
• 검색 중심 사전 훈련에서 비대칭 모델 구조와 마스킹 비율의 효과를 탐색하기 위해.

제안 방법

• RetroMAE는 입력 문장을 인코더와 디코더에 대해 다르게 마스킹하는 새로운 마스크된 오토인코딩 워크플로우를 사용한다.
• 인코더는 중간 수준의 마스킹(15–30% 마스킹)을 적용한 입력에서 문장 임베딩을 생성하는 전체 크기의 BERT이다.
• 디코더는 단일 레이어 트랜스포머로, 문장 임베딩과 강력하게 마스킹된 입력을 사용하여 원래 문장을 복원한다(50–70% 마스킹).
• 두 개의 스트림 어텐션과 위치별로 특화된 어텐션 마스크를 도입하여 디코딩을 향상시켜, 모든 토큰이 고유한 맥락을 활용하여 복원되도록 보장한다.
• 비대칭 마스킹 비율을 사용하여 디코더의 복원 난이도를 높이면서도 인코더의 입력 품질을 유지한다.
• 모델은 신호 활용도와 표현 품질을 극대화하는 데 중점을 두고 위키피디아, BookCorpus, MS MARCO에서 사전 훈련된다.

실험 결과

연구 질문

• RQ1 비대칭 마스킹 비율을 사용하는 마스크된 오토인코딩 프레임워크가 밀도 있는 검색을 위한 문장 표현 품질을 향상시킬 수 있는가?
• RQ2 강화된 디코딩을 적용한 단일 레이어 디코더의 사용이 데이터 효율성과 검색 성능에 어떤 영향을 미치는가?
• RQ3 검색 중심 사전 훈련 설정에서 인코더와 디코더에 적합한 최적의 마스킹 비율은 무엇인가?
• RQ4 비대칭 모델 아키텍처(전체 인코더, 경량 디코더)가 대칭 설계보다 더 높은 성능을 내는가?
• RQ5 강화된 디코딩이 사전 훈련 단계에서 훈련 신호의 다양성과 모델 일반화 능력을 얼마나 향상시키는가?

주요 결과

• RetroMAE는 BEIR 벤치마크에서 제로샷 평균 점수 45.2를 기록하여 기존 방법들을 크게 앞서며 성능을 뛰어넘었다.
• DPR에서의 지도 학습 미세조정 시, RetroMAE는 여러 밀도 있는 검색 벤치마크에서 최신 기술 수준의 성능을 달성했다.
• 두 개의 스트림 어텐션과 위치 기반 마스크를 사용한 강화된 디코딩은 데이터 효율성을 향상시키며 뚜렷한 성능 향상을 이끌었다.
• 일반적으로 더 깊은 디코더보다 한 레이어 디코더가 더 높은 성능을 보였으며, 더 큰 디코더를 사용할 경우 강화된 디코딩을 비활성화해야 하며 이는 성능 저하를 초래했다.
• 인코더 마스킹 비율을 높이면(최대 30%) 및 디코더 마스킹을 강하게 적용하면(50–70%) 검색 품질이 향상되었으며, 디코더 마스킹 비율이 0.5일 때 최적의 성능을 기록했다.
• 제거 실험을 통해 비대칭 마스킹과 강화된 디코딩이 훈련 신호 활용도와 모델 효과성을 극대화하는 데 핵심 요소임을 확인했다.