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

 

 

Abstract 

LLM은 파라미터에 지식을 저장하고, 다운스트림 NLP task에 미세조정할 때 SOTA를 달성하는 좋은 성능을 보인다. 그러나 정확하게 지식에 접근하고 조작하는 성능은 제한되어 있다. 또한 출력한 정보의 출처를 제공하고 새로운 지식 체계의 정보를 업데이트하는 것은 해결되지 않은 문제이다. 논문은 pre-trained parametirc memory와 non-prametric memory를 결합한 언어 생성 모델 Retrieval-Augmented Generation(RAG)를 제안한다. RAG의 parametric memory는 사전학습된 seq2seq 모델이고 non-parametric memory는 사전학습된 neural retriever로 위키피디아 dense vector index에 접근한다. 언어 생성 task에서 RAG는 parametric-only seq2seq baseline 보다 더 구체적이고 사실에 근거한 언어를 생성할 수 있다.

 

1 Introduction

사전학습 언어 모델은 데이터에서 상당한 양의 지식을 학습한다. 외부 지식에 접근 없이 파라미터화된 내부 지식 베이스를 활용할 수 있다. 한편 이러한 발전은 기억하고 있는 정보를 쉽게 확장하거나 수정할 수 없고 예측에 대한 인사이트를 제시할 수 없다. 또한 환각(hallucinations)을 유발하는 단점이 있다.

Hybrid 모델은 parametric memory와 retrieval 기반의 non-parametric memories를 결합했다. Hybrid 모델은 저장된 지식을 직접 수정하거나, 확장할 수 있고, 직접 지식에 접근해 검사하고 해석할 수 있다. 논문은 parametric memory는 seq2seq transformer로 non-parametric memory는  dense vector index로 변환한 Wikipedia에 접근할 수 있는 사전학습된 neural retriever로 RAG를 구성했다. Retriever(Dense Passage Retriever-DPR)는 조건으로 주어진 입력으로 잠재 문서(latent documents)를 제공한다. 그리고 seq2seq 모델(BART)이 잠재 문서와 input을 함께 조건으로 입력받고, 출력을 생성한다. 논문은 동일한 문서가 모든 토큰을 담당하면 per-output, 다른 문서가 다른 토큰을 담당하면 per-token을 기준으로 top-k approximation으로 latent documents를 제외한다. RAG는 seq2seq task에서 미세조정할 수 있으므로 generator와 retriever 모두 학습한다. 또한 Pre-trained access mechanism으로 knowledge에 access 할 수 있고 추가적인 학습이 필요 없다.

RAG는 parametric과 non-parametric memory의 결합으로 knowledge-intensive tasks에서 이점이 있음을 강조한다. RAG는 사실적이고 구체적인 응답이 가능하고, 세상이 변화함에 따라서 모델의 지식을 업데이트하기 위해 non-parametric memory를 업데이트할 수 있다.

 

2 Methods

RAG는 input sequence $x$, 검색된 텍스트 문서 $z$, 이 둘을 추가 컨텍스트로 사용하여 target sequence $y$를 생성한다.

RAG는 두 가지로 구성된다.

1. retriever
   • $p_\eta(z|x)$
   • $\eta$: 주어진 $x$로 반환한  top-K text passages
2. generator 
   
   • $p_\theta(y_i|x,z,y_{1:i-1})$
     1. input $x$
     2. $\theta$ 생성한 current token
     3. $x$로 검색된 passage $z$
     4. $x$와 $z$로 생성한 현재까지 토큰$y_{1:i-1}$
     5. $x, z, y_{1:i-1}$를 input으로 현재 토큰을 생성

retriever와 generator를 end-to-end로 학습하기 위해서 검색된 문서를 잠재변수(latent variable)로 처리한다. 논문은 generated text의 distribution을 생성하기 위해서 다른 두 가지 모델로 latent document를 주변화한다(marginalize over).  한 가지는 RAG-Sequence이다. 모델은 같은 document를 사용해 각 target token을 예측한다. 다른 하나는 RAG-Token이다. 각 토큰은 다른 document를 참조한다.

 

2.1 Models

RAG-Sequence Model

RAG-Sequenece 모델은 전체 sequence를 생성하기 위해서 검색된 문서를 사용한다. 검색된 document는 하나의 latent variable로 다루어진다. latent variable은 top-K approximation을 통해 seq2seq  확률 $p(y|x)$를 얻기 위해 주변화된다. retriever는 top K documents를 검색하고, generator는 각 document에 대한 sequence 확률을 출력한다.

 

 

RAG-Token Model

Rag-Token 모델은 target token마다 다른 latent document를 이끌어내고 주변화한다. 따라서 generator는 답변을 생성할 때 여러 documents를 참조한다. 위의 공식의 순서는 다음과 같다.

1. retriever가 top K documenrts를 검색한다.
2. generator가 각 문서의 다음 output token의 분포를 생성한다.
3. output token을 사용해 과정을 반복한다.

 

2.2 Retriever: DPR

• 리트리버는 DPR로 기반이다.
• DPR은 bi-encoder architecture를 따른다.
• $d(z)$는 BERT 기반 document encoder로 생성한 dense representation이다.
• $q(x)$는 BERT 기반 query encoder로 생성한 query representation
• $top-k(p_{\eta}(\centerdot |x))$는 k
• 사전 확률(prior probability) $p_{\eta}(z|x)$ 는 document list $z$ 중에서 가장 높은 $k$ 개의document를 계산하며, 이는 MIPS(Maximum Inner Product Search) 문제로 sun-linear time으로 풀 수 있다. 
•  DPR에서 사전학습된 bi-encoder를 사용해 retriever를 initialize하고, document index를 빌드하여 학습한다.
• 논문은 document index를 non-parametric memory로 언급한다.

 

2.3 Generator: BART

• $p_{\theta}(y_i | x, z, y_{1:i 1})$
• encoder-decoder BART-large 모델로 사전학습된 400M 파라미터의 seq2seq transformer를 사용한다.
• input $x$와 검색된 $z$ 를 concatenate하여 BART로 생성한다.
• 논문은 BART generator parameters $\theta$를 parametric memory라고 한다.

 

2.4 Training

$(x_j, y_j)$ 가 fine-tuning 데이터로 주어졌을 때 타겟에 negative marginal log-likelihood $\sum_j - \log p(y_j|x_j)$ 를 Adam으로 최소화한다. document encoder $BERT_d$를 훈련을 통해 업데이트하면 주기적으로 document index를 업데이트해야 하므로 비용이 많이든다. 논문은 document 학습이 불필요하다 판단하고, document encoder(index)를 fix했다. query encoder인 $BERT_q$와 $BART$ generator만 fine-tuning했다.

 

 

2.5 Decoding

test에서 RAG-Sequence와 RAG-Token은 서로 다른 방법으로 $argmax_y p(yx)$ 근사가 필요하다.  

 

RAG-Token

RAG-Token은 표준 모델로 autoregressive seq2seq generator이다. decorde하기 위해서 $p\prime_{\theta}(y_i|x, y_{1:i-1}) $를 standard beam decorder에 연결한다.

 

RAG-Sequence

document $z$를 각각 beam search 실행하고, $p_{\theta}(y_i|x, y_{1:i-1})$ 식을 이용해 각각 점수를 준다. 각각의 documents에서 beam search 수만큼 문장을 생성하게 된다. 예를 들면 2개의 passage로 beam search 3개의 문장을 생성하면, 1개의 passage에서 3문장씩 총 6문장을 생성한다. 이후 생성된 문장을 각 passage에서 확률을 더해 점수가 높은 문장을 구한다.

문제는 모든 passage에서 beam search로 문장을 생성하지 않을 때이다. 1번 passage에서 "나는 점심을 먹어"라는 문장을 생성하고 2번 passage에서 생성하지 않았을 때,  "나는 점심을 먹어" 문장의 확률을 1번 passage는 구할 수 있지만 2번 passage가 생성할 확률을 알 수 없다. 이때 " Thorough Decoding"은 forward를 통해서 2번 passage를 이용했을 때 logit을 다시 구한다. 단점은 추가적인 연산이 필요하다. 논문은 “Fast Decoding”으로 등장하지 않은 passage의 확률은 0으로 변경하여 추가 연산 없이 확률을 구했다.

 

Experiments and Result

Wikipedia article 21M document를 100word로 chunking 후 dump했다. FAISS를 사용해 MIPS index를 빌드했다. Open-domain Question Answering의 4개의 task에서 SOTA를 달성했다. DPR’s retriever로 initialize하여 pretrain 없이 task에 지도학습하였다. 또한 다른 모델은 DPR로 검색한 모델에 답이 없으면 정답률이 0%이지만 RAG는 retriver로 정답이 없는 document 를 가져왔을 때도 11.8% 정답을 맞출 수 있었다.