[논문 리뷰] Fine-tune the Entire RAG Architecture (including DPR retriever) for Question-Answering

Shamane Siriwardhana, Rivindu Weerasekera외 2명|arXiv (Cornell University)|2021. 06. 22.

Topic Modeling참고 문헌 7인용 수 4

한 줄 요약

이 논문은 질문 인코더와 문서 인코더를 공동으로 훈련함으로써, 밀도 있는 문단 검색(DPR) 검색기까지 포함된 전면적 Retrieval-Augmented Generation(RAG) 아키텍처의 엔드 투 엔드 미세조정을 제안한다. 이 방법은 병렬 처리를 사용하여 훈련 중 지식 기반의 재인코딩 및 재색인화와 같은 기술적 과제를 해결하며, 원래의 RAG 대비 SQuAD에서 정확도가 12% 향상되어 정확히 일치하는 정확도가 40.02%에서 28.12%로 향상된다.

ABSTRACT

In this paper, we illustrate how to fine-tune the entire Retrieval Augment Generation (RAG) architecture in an end-to-end manner. We highlighted the main engineering challenges that needed to be addressed to achieve this objective. We also compare how end-to-end RAG architecture outperforms the original RAG architecture for the task of question answering. We have open-sourced our implementation in the HuggingFace Transformers library.

연구 동기 및 목표

• 기존 Hugging Face 구현에서 지원되지 않았던, DPR 검색기를 포함한 전체 RAG 아키텍처의 엔드 투 엔드 훈련을 가능하게 하기 위해.
• 지식 기반 임베딩 갱신 및 재색인화 과정에서 발생하는 기술적 과제를 해결하기 위해, 이는 계산 비용이 많이 들고 시간이 오래 걸리는 작업이다.
• 질문 인코더와 문서 인코더를 공동으로 미세조정하는 것이 동결된 문서 인코더 대비 질문-답변 성능을 향상시키는지 평가하기 위해.
• 더 넓은 연구 사용과 재현 가능성을 위해 구현을 오픈소스화하기 위해.

제안 방법

• Hugging Face Transformers 라이브러리를 확장하여, DPR 문서 인코더까지 포함된 RAG 모델의 엔드 투 엔드 훈련을 지원하도록 구현하였다.
• 이중 BERT 아키텍처를 사용: 하나는 질문을 인코딩하고, 다른 하나는 문서를 인코딩하며, 둘 다 공동으로 미세조정되었다.
• 파이썬의 멀티프로세싱과 HuggingFace Datasets를 사용하여 지식 기반의 병렬 재인코딩 및 재색인화를 구현하여 메인 훈련 루프를 차단하지 않도록 하였다.
• 효율적인 근접 이웃 검색을 위해 FAISS를 활용하고, 각 재인코딩 단계 이후 벡터 색인을 갱신했다.
• 문서 검색 확률 p(z|x)를 통해 기울기 역전파를 적용하여 질문 인코더와 문서 인코더 양쪽을 동시에 업데이트하였다.
• 성능과 계산 비용의 균형을 맞추기 위해 주기적인 업데이트 전략—훈련 단계마다 N단위 간격으로 재인코딩 및 재색인화를 수행—을 적용하였다.

실험 결과

연구 질문

• RQ1 RAG 내의 DPR 검색기를 엔드 투 엔드로 미세조정하는 것이 동결된 문서 인코더 대비 질문-답변 성능 향상에 기여하는가?
• RQ2 RAG 훈련 중에 동적으로 업데이트되는 지식 기반을 유지하는 데 있어 핵심적인 기술적 과제는 무엇인가?
• RQ3 지식 기반의 재인코딩 및 재색인화 빈도는 모델 성능과 훈련 효율성에 어떤 영향을 미치는가?
• RQ4 질문 인코더와 문서 인코더를 공동으로 훈련하는 것이 도메인 적응 및 일반화 능력을 향상시키는가?

주요 결과

• DPR 검색기를 포함한 RAG 아키텍처의 엔드 투 엔드 미세조정으로, SQuAD 데이터셋에서 정확히 일치하는 정확도가 원래 RAG의 28.12%에서 40.02%로 향상되었다.
• 성능 향상은 질문 및 문서 인코더의 공동 최적화가 검색 및 생성 품질을 향상시킨다는 것을 보여준다.
• 훈련 중 지식 기반의 재인코딩 및 재색인화는 계산 비용이 매우 높지만, 병렬 처리 덕분에 합리적인 훈련 시간 내에 구현이 가능해졌다.
• 연구 결과는 훈련 중 문서 인코더를 갱신함으로써 쿼리와 관련된 문서 간의 정렬이 향상되며, 특히 위키백과 기반 외의 도메인 특화 환경에서 뚜렷한 효과를 보임을 확인하였다.
• 오픈소스화된 구현은 Hugging Face Transformers와 PyTorch Lightning를 사용하여 연구자들이 해당 방법을 재현하고 확장할 수 있도록 한다.