[NLP] RAG: Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

RAG: Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

 

link : https://arxiv.org/pdf/2005.11401

 

 

✲ Abstract

 

기존 대량의 코퍼스에 있는 factual knowledge를 가지고 대형 사전 훈련 언어 모델은 down-stream NLP tasks에 파인튜닝되어 자연어 생성과 이해에서 좋은 성능을 보여주었다. 하지만 지식 집약적인 tasks의 성능은 뒤쳐져 있는 편이며 새로운 지식에 접근하거나 정교하게 조작하는 것은 한계가 있다. 이를 위해 해당 논문에서는 RAG (Retrieval-Augmented Generation)라는 새로운 Fine-tuning recipe를 제안한다. RAG는 다운 스트림 task 뿐 아니라 general-purpose를 위해 제안되었으며 pre-trained parametric 모델에 non-parametric memory를 통합한다. 먼저 seq2seq로 사전 학습된 parametric model이 있을 때 위키피디아 등 새로운 지식의 dense vector index가 non-parametric memory가 되어 parametric model에 접근하여 새로운 정보를 통해 답변을 생성할 수 있다.

 

 

✲ Introduction

 

Pre-trained neural language model는 parametrized implicit knowledge base로써 학습 된 후 외부 메모리 접근 없이 언어 생성이 가능하다. 하지만 LM은 쉽게 해당 메모리를 수정할 수도 확장할 수도 없다. 또한 prediction에 대해 간단하게 인사이트를 제공할 수도 없으며 문장 자체에는 문법적인 오류가 없어 그럴듯하지만 허구이거나 거짓인 답변, "hallucinations"를 생성한다는 문제점도 존재한다.

 

이를 위해 해당 논문에서는 prametric-menory generation models와 non-parametric memory를 사용해 일반적인 목적에 맞는 파인 튜닝 기법을 소개한다. end-to-end로 학습된 확률 모델에서 이러한 컴포넌트들을 결합한다.(아래의 figure 1) Dense passage retriever는 입력에 대해 조건화된 latent document를 제공하고 seq2seq 모델은 입력과 함께 이러한 latent document에 대한 조건을 지정하여 출력을 생성한다. per-output basis (동일한 문서가 모든 토큰을 담당) 또는 per-token basis (다른 문서가 서로 다른 토큰을 담당) 로 top-k approximation을 통해 latent document를 marginalize한다. 다른 모델들과 마찬가지로 RAG는 seq2seq로 파인튜닝 되기 때문에 sequence generator와 retriever가 함께 학습된다. 여기서 제안하는 pre-trained access mechanism을 사용하면 이후 추가적인 학습 없이도 지식에 접근할 수 있는 능력을 모델이 갖게 된다. 마치 우리가 지식 집약적인 작업을 할 때 단순히 가지고 있는 지식 외에 외부 정보를 이용하는 것과 같다.

 

 

✲ Methods

 

 

위 그림에서 볼 수 있듯 RAG는 두 개의 컴포넌트를 이용한다.

 

1. Retriever ( p_\eta(z|x) )
: query x가 주어졌을 때 텍스트 구절에 대한 분포(top-K가 잘린)를 반환함

2. Generator ( p_\theta(y_i | x, z, y_{1 : i-1} )
: 원래 인풋 x와 검색된 구(retrieved passage) z, 이전까지 생성된 토큰 y_{1:i-1}이 주어졌을 때 현재 토큰을 생성함

 

 

‣ Models

 

검색기(Retriever), 생성기(Generator)를 end-to-end로 학습하기 위해 검색된 문서를 latent variable로 취급한다. 그리고 생성된 텍스트의 분포를 만들기 위해 latent documents를 marginalize하는 두 가지 기법을 소개한다.

 

1. RAG-Sequence : 모델은 각 타겟 토큰을 예측하기 위해 같은 문서를 사용

2. RAG-Token : 모델은 각 타겟 토큰을 예측하기 위해 다른 문서를 사용

 

 

검색기를 사용해 검색된 K개의 문서를 입력으로 받은 generator가 답변을 생성하는 것은 같다.

RAG-sequence는 검색된 하나의 동일한 문서에서 전체 답변을 생성하고 RAG-Token은 각 토큰에 따라 서로 다른 문서를 사용할 수 있다. 

 

 

‣ Retriever : DPR

 

검색기는 QA의 검색 구성요소를 개선하는 이전 연구의 방법 DPR(Dense Passage Searcher) 을 그대로 사용한다. 먼저 문서 D가 있는 경우 문서를 같은 길이를 가지는 텍스트 구로 쪼개서 basic retrieval units로 만든다. 총 M개의 passage 가 만들어지면 코퍼스 C = {p_1, p_2, .. p_M}이 된다. question 또는 query가 주어지면 passages 중 하나를 찾아 답변을 위한 passage안의 문장 범위를 찾는다. Retriever는 저차원 연속 공간 안에서 모든 구절을 색인화하여 쿼리와 관련된 상위 k 구절을 효율적으로 검색할 수 있도록 하는 것이다.

 

 

쿼리와 텍스트 구를 위한 인코더가 각각 있고 여기서는 BERT를 사용했다.

BERT_d는 메모리에 저장된 documents z를 입력으로 받고 BERT_q는 쿼리 텍스트 x를 입력으로 받아 각각 d-dimensional 실수 벡터를 생성한다. 그리고 검색기는 exp dot product로 이 둘의 유사도를 가지고 쿼리와 관련있는 텍스트 구를 검색하게 된다.

 

상위 관련있는 k개의 문서를 계산하는 것은 MIPS(Maximum Inner Product Search) 문제이며 이는 대략 sub-linear 타임으로 계산 가능하다 (참고). 

 

 

‣ Generator : BART

 

generator는 어떤 인코더-디코더 구조의 모델도 사용 가능하다. 해당 논문에서는 BART-large 모델을 사용했으며 약 400M 파라미터를 가진 모델이다. 입력 x와 검색된 content (문서에서 쪼개진 text passage p_i 중 특정 부분 문장들) z는 간단하게 concat 하여 입력으로 넣는다.

 

 

‣ Training

 

학습은 어떤 문서를 검색해야 하는지 직접적인 supervision을 주지 않는다. 단지 파인튜닝을 위한 코퍼스의 입력-출력 페어 (x_j, y_j)가 주어지면 Adam SGD를 사용해 각 타겟의 negative marginal log-likelihood  -\sum_j log p(y_j | x_j)를 미니마이즈 한다.

그리고 문서 인코더 BERT_d를 학습하는 동안 업데이트 하는 것은 문서 인덱스를 주기적으로 업데이트해야하므로 비용이 많이 든다. 따라서 BERT_d는 고정하고, 쿼리 인코더 BERT_q, Generator만 학습한다.

 

 

‣ Decoding

 

테스트 타임에 RAG-sequence와 RAG-token은 토큰 생성 argmax_y p(y|x)를 근사 하기 위해 서로 다른 방법이 필요하다.

 

 

 

 

RAG-token의 경우 token별로 marginal 하기 때문에 모델은 일반적인 autoregressive seq2seq generator가 토큰을 생성하는 transition probability와 동일하게 토큰을 생성한다. retriever p_\eta와 generator p_\theta의 곱을 top k개의 문서 z에 대해 summation한 것이고 위의 p'_\theta이다. p'_\theta를 가지고 beam search하여 토큰을 생성한다.

 

 

 

 

 

RAG-sequence의 경우 문장 끝까지 생성한 후 document에 대해 marginalize한다. 때문에 모델에서 likelihood p(y|x)는 per-token liklihood로 쪼개지지 않아 single beam search로 디코딩 할 수 없다.

대신, 각 문서 z에 대해 generator p_\theta(y_i | x, z, y_{1:i-1}) 스코어로 beam search를 수행할 수 있다. 하지만 hypotheses Y 중 일부는 모든 문서의 빔에 나타나지 않을 수 있다. 예를 들어 top_k=3일 경우 y_1은 z_1, z_2에는 있지만 z_3에는 없을 수 있다. 그렇게 되면 p(y_1|x, z_3)는 구할 수 없다. 

 

 

출처 : https://gbdai.tistory.com/67

 

모든 빔에 대한 Y의 확률을 추정하기 위해 y beam에 나타나지 않는 각 document z에 대해 추가적인 forward pass를 수행해야 한다. 먼저 각 document z별로 beam search를 수행해 p(y|x,z)를 구한다. 그리고 이때 등장하지 않은 p(y|x,z)를 위해 추가적인 forward pass (model에 값을 넣어 p_\eta를 구함)를 수행해 p_\eta(z|x)를 구해 marginalize를 하여 p(y|x,z)를 구하고 이를 합산한다. 이를 Thorough Decoding이라 한다.

 

출처 : https://gbdai.tistory.com/67

 

하지만 시퀀스가 긴 경우 |Y|도 커질 것이므로 많은 forward pass가 필요하여 좀 더 효율적인 디코딩 방법을 적용해야 한다. 빔 서치 동안 특정 z에 대해 등장하지 않은 y에 대한 p(y|x,z)를 0으로 처리한다. 이는 candidate set Y가 생성된 후 추가적인 forward pass는 필요 없다. 이를 Fast Decoding이라 한다.

 

 

 

---

 

‣ Review

 

RAG를 간단하게 요약하면 모델에게 주고 싶은 document가 있으면 이를 각 text passage로 쪼개고 벡터 인덱스로 구성한다. query 벡터와 text passage의 유사도를 통해 관련된 top k개의 문서를 추출하고 쿼리와 합쳐 generator에 입력으로 만든다. 기존에 쿼리와 이전까지 생성된 타겟 토큰만을 입력으로 받아 next token을 생성했다면 거기에 검색된 문서가 추가되는 것이다. 내용이나 알고리즘은 간단하고 직관적이지만 문서를 벡터화, 색인화하고 유사한 부분을 검색하는 것을 빠르게 하기 위한 엔지니어링이 중요해 보인다.

 

2022년까지의 데이터로 학습시킨 언어 모델이 있을 때 이후 데이터에 대해 업데이트 하기 위해서는 추가적인 학습이 또 필요하므로 충분한 양의 데이터가 모일 때까지 모델을 추가 학습시키기는 어렵다. 따라서 모델에 새로운 정보를 계속 업데이트 해야한다면 다른 방법이 필요하다. 이때 RAG를 사용할 수 있으며 일반적인 파인튜닝 없이 쿼리 인코더와 Generator를 한번 학습하면 새로운 지식의 vector index만 구성함으로써 모델에 새로운 정보를 줄 수 있다는 것이 큰 장점이다. RAG관련 프레임워크도 쉽게 사용 가능해서 활용도가 높아 보인다. 특히 대형 언어 모델일 수록 언어 모델이 생성하는 답변을 원하는 방식으로 조정하는 것이 어렵고 모델의 생성을 세세하게 조정할 수 없다면 실제 어플리케이션에 이용하기 힘들다. 근본적인 방법은 아니겠지만 응용단에서 RAG를 이용하면 이러한 부분이 조금 해결될 수 있는 것 같다.