본문 바로가기

논문 리뷰

Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks 논문 리뷰

Lewis, Patrick, et al. "Retrieval-augmented generation for knowledge-intensive nlp tasks." Advances in Neural Information Processing Systems 33 (2020): 9459-9474. 

 

Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Requests for name changes in the electronic proceedings will be accepted with no questions asked. However name changes may cause bibliographic tracking issues. Authors are asked to consider this carefully and discuss it with their co-authors prior to reque

proceedings.neurips.cc

 

최근 LLM모델을 이용하기 위한 방법론 중 하나로 많이 사용되고 있는 RAG 논문이다.

 

Abstract

사전훈련된 LLM들은 자신들의 파라미터에 실제로 지식을 저장하는 능력을 보여주었고, 다른 NLP task에 fine-tuned될 때도 높은 성적을 보였다. 하지만 여전히 제한적인 능력을 보이며 특정 작업에 중심적인 아키텍처들에 비해 부족한 성능을 보인다. 또한 학습된 지식을 업데이트하는것은 여전히 연구과제로 남아있다. 우리는 사전학습된 paramenter와 non-parameter을 결합한 RAG(Retrieval-Augmented Generation)을 위한 fine-tuning recipe을 탐구한다. 여기서 사전학습된 parameter란 사전학습된 seq2seq모델을, non-parameter은 사전 학습된 neural retriever을 사용하여 접근한 위키백과의 dense vector index을 뜻한다.

 

우리는 두가지의 RAG을 비교한다. 첫번째는 생성된 전체 시퀀스에 대해 검색된 하나의 동일한 문서를 사용하는 것이고 두번째는 각 토큰마다 다른 문서를 사용하는 방식이다. 우리는 세가지 open-domain QA와 언어 생성의 작업에서 다른 모델보다 뛰어난 성능을 보였다.

 

 

Methods

 

우린 RAG 모델을 탐구하는데 대해 입력 시퀀스 x를 사용하여 문서 z를 검색하고, 이를 타겟 시퀀스 y를 생성하기 위한 추가적인 요소로 활용했다.

우리 모델은 Figure 1에 있는것 처럼

검색기 \(P_\eta\) : 매개변수 \(\eta\)로 구성된, 질의 x에 대해 top-K 로 잘려진 문서분포를 반환한다.

생성기 \(P_\theta\) : 매개변수 \(\theta\)로 구성된, 이전 i-1개의 토큰, 원래 입력 x, 검색된 문서 z을 기반으로 현재 토큰 \(y_i\)을 생성한다.

 

우리는 검색기와 생성기의 end-to-end 학습을 위해 검색된 문서는 latent variable로 처리된다. 우리는 output y을 생성하기 위해 latent variable된 문서 z을 제공하기 위한 두가지 다른 모델을 제안한다.

1. RAG-Sequence : 한 번 검색된 문서를 z을 이용해 y를 생성할 때 계속 사용한다.

2. RAG-Token : 각 토큰을 생성할 때마다 필요한 문서 z를 계속해서 선택한다.

 

RAG-Sequence Model 

하나의 검색된 문서를 사용하여 전체 출력 시퀀스를 생성한다. 기술적으로는, 검색된 문서를 하나의 잠재 변수로 간주하며, 이를 marginalization하여 seq2seq 확률 을 계산한다. 이를 위해 Top-K개의 문서를 검색기로 가져오고, 각 문서에 대해 생성기가 출력 시퀀스의 확률을 계산한다. 이후, 이 확률들을 결합하여 최종적으로 marginalization된 확률을 얻는다.

 

즉, 검색된 여러 문서를 활용하되, 최종 출력은 한 문서를 기준으로 생성된 결과를 기반으로 marginalization 과정을 거쳐 확률 분포를 산출하는 방식이다.

 

 RAG-Token Model

각 토큰마다 다른 잠재 문서를 선택한다. 구체적으로는 Top-K개의 문서를 검색기로 가져온다. 생성기는 각 문서에 대해 현재 출력 토큰의 확률 분포를 계산한다. 이를 기반으로 marginalization하여 최종 확률 분포를 얻는다. 이후 다음 출력 토큰에 대해 동일한 과정을 반복한다.

이 방식은 각 단어를 생성할 때마다 다른 문서를 참고할 수 있어, 더 다양한 정보와 내용을 반영한 답변을 생성할 수 있다는 장점이 있다.

 

우리의 모델에서 검색기는 DPR 모델을, 생성기는 BART 모델을 사용하였다.

 

Training

RAG 모델은 검색기와 생성기가 어떤 문서를 검색해야 되는지 supervision없이 동시에 학습된다. 모델은 negative marginal log-likelihood를 최소화 하는 방향으로 훈련되며 SGD, Adam을 사용했다. 또한 학습 중 \(document encoder BERT_d\)을 업데이트 하는 것은 큰 비용이 소모되지만 필수적인 요소는 아닌 것을 발견했다.

따라서, 1. document encoder, document index은 고정하고, 2. \(query encoder BERT_q\), 생성기 BART만 미세조정한다.

 

 

Experiment

우리는 다양한 실험에서 하나의 위키백과 꾸러미만을 non-parametric 자료로 사용했다(December 2018 dump). 각 문서들을 100단어 단위로 나누어 서로 겹치지 않게 분할하여 총 2100만개의 문서가 생성된다. Document encoder을 사용해 문서들끼리의 embedding을 계산하고, 이를 기반으로 MIPS 인덱스를 생성한다. 이 과정에서 FAISS를 활용하며, 검색 속도를 높이기 위해 HNSW 근사 방식을 사용한다.

 

MIPS (Maximum Inner Product Search)

MIPS는 벡터 간 **내적(inner product)**을 기준으로 가장 큰 값을 가지는 벡터(즉, 가장 유사한 항목)를 빠르게 찾는 문제를 해결합니다.

  • 내적의 의미: 두 벡터가 유사할수록 내적 값이 큽니다. 따라서 MIPS는 벡터의 유사도를 계산하고, 가장 관련성이 높은 결과를 반환합니다.

FAISS (Facebook AI Similarity Search)

FAISS는 대규모 벡터 데이터베이스에서 효율적으로 유사 항목을 검색하기 위한 라이브러리입니다.

 

HNSW (Hierarchical Navigable Small World):

  • 효율적인 근사 최근접 이웃 검색(ANN: Approximate Nearest Neighbor) 알고리즘.
  • 그래프 구조를 사용해 탐색 경로를 최적화하여 빠르게 유사 항목을 찾음.
  • 계층적 설계 덕분에 높은 검색 정확도를 유지하면서도 속도를 크게 개선.

MIPS + FAISS의 결합 효과

RAG 모델에서 FAISS와 MIPS를 사용하면, 수백만 개의 문서(위키백과 청크) 중에서 사용자의 쿼리에 가장 적합한 문서를 빠르게 검색할 수 있습니다.

  • 문서 임베딩: 문서를 벡터화해 FAISS에 저장.
  • 쿼리 임베딩: 사용자의 입력을 벡터화.
  • MIPS 검색: 쿼리 벡터와 문서 벡터 간 내적 계산으로 가장 유사한 문서를 찾아냄.

 

우리는 각 쿼리에 대해 top-k, \(k\in{5,10}\) 문서를 검색한다.

 

이후 자세한 실험결과는 논문을 참조하길 바람