[논문 리뷰] RAG, Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Posted Jul 25, 2024

By Luna Lee

21 min read

이번 글에서 리뷰할 논문은 최근 LLM에서 많이 사용되는 RAG(검색 증강 기법)의 개념이 처음으로 등장한 논문이다. 이 논문에서는 Parametric 메모리의 모델(Generator)을 fine-tuning할 때 non-parametric 메모리(Retriever)를 활용함으로서, 모델의 지식을 업데이트할 수 있음을 보여준다. 다만 Retriever과 Generator를 fine-tuning한다는 점이 현재 LLM에서 많이 사용되는 RAG와는 차이점이 있다. LLM이 많이 사용되기 이전의 논문이라 fine-tuning으로 모델을 업데이트 했지만, LLM에서의 RAG는 모델 업데이트 없이 진행된다. 차이점에 유의해서 이해해보자.

Facebook AI Research NIPS 2020

Paper: https://arxiv.org/abs/2005.11401
Demo: https://huggingface.co/facebook/rag-token-nq

📖 핵심 훑어보기 !!

Parametric 메모리(Generator)와 non-parametric 메모리(Retriever)를 결합하는 검색 증강 생성(RAG) 방식으로 학습 진행
Retriever(DPR)는 입력에 따라 latent document를 검색하고, Generator(BART)는 입력과 함께 검색된 latent document를 사용하여 target output 생성.
Latent documents들을 distribution으로 변환하기 위하여 marginalize를 수행하는데, marginalize 방식에 따라 두 가지 모델(RAG-Sequence과 RAG-Token 모델) 제안

Introduction

Pretrain된 language 모델은 방대한 데이터에 대한 심층적인 지식을 학습하고, 모델의 parameter를 통해 학습된 암묵적 knowledge를 기반으로 외부 메모리(e.g. 특정 지식,정보가 저장되어 있는 메모리)에 접근하지 않고도 정보를 제공한다. 이러한 모델은 암묵적인 지식에 기반하기 때문에 메모리를 쉽게 확장하거나 수정하기 어렵다. 모델이 학습하지 않은 추가적인 정보를 주입한다거나, 이전에 학습한 일부의 지식을 변경하는 작업이 어렵고, 흔히 말하는 “hallucinations” 현상, 즉 잘못된 정보를 제공하는 일이 생긴다.

본 논문에서는 parametric 메모리와 non-parametric 메모리(즉, 검색 기반)를 결합하는 검색 증강 생성(RAG) 방식으로 학습을 진행했다. 좀 더 자세히는 parametric 메모리가 pre-train된 seq2seq transformer이고, non-parametric 메모리는 Wikipedia의 dense vector index라고 볼 수 있는데, 이 index에는 pre-train된 neural retriever를 통해 접근한다. Retriever(Dense Passage Retriever, DPR)는 입력에 따라 조건화된 latent document를 제공하고 seq2seq 모델(BART), 즉 Generator는 입력과 함께 이러한 latent document를 지정하여 출력을 생성하는 방식이다. 이 때 Retriever과 Generator는 동시에 학습된다.

Parametric 메모리의 모델을 fine-tuning할 때 이렇게 non-parametric 메모리를 활용함으로서, non-parametric 메모리를 대체하여 모델의 지식을 업데이트할 수 있음을 보여준다. 또한 모델이 더 사실적이고 구체적이며 다양한 응답을 생성할 수 있음을 증명했다.

Method

RAG 모델은 위 그림과 같이 input sequence $x$를 사용하여 text documents $z$를 검색하고, 이를 추가 context로 활용하여 target sequence $y$를 생성한다. 모델은 두 가지로 구성되어있다.

Retriever $p_{\eta}(z∣x)$ with parameter $\eta$: query $x$가 주어지면 해당 텍스트에 대한 (top-K)distribution을 반환함.
Generator $p_\theta(y_i ∣ x, z, y_{1:i-1})$ parametrized by $\theta$: 이전의 $i-1$ token $y_{1:i-1}$, original input $x$, 검색된 구절 $z$를 기반으로 현재 token 생성.

Retriever와 Generator를 동시에 학습시키기 위해 검색된 documents들을 latent variable로 취급한다. 이 latent documents들을 distribution으로 변환하기 위하여 각각 다른 marginalize**를 사용하는 두 가지 모델을 제안했다.

RAG-Sequence: 동일한 document를 사용하여 각각의 target token을 예측
RAG-Token: 다른 document를 기반으로 각각의 target token을 예측

Marginalize**
Joint distribution 즉, 두 개의 사건 $X, Y$가 동시에 일어날 확률이 주어졌을 때, 하나의 변수 $Y$에만 관심이 있고 변수 $X$에는 관심이 없는 상황을 가정해보자. 이 때, 변수 $Y$에 대한 확률값을 추정하기 위해 $X$의 값에 대해 확률을 합산(적분)하여 $Y$의 확률만 남기는 것을 의미한다.
여기서 target $y$에 대한 확률만을 구하기 위해서 latent document $z$에 대한 확률을 marginalize 한다고 볼 수 있다. 입력 $x$ → latent document $z$ → target sequence $y$를 생성하는 과정에서 target $y$에 대해 supervised 학습을 진행하지만, document $z$를 구하는 과정에는 따로 직접적인 정답이 주어지지 않고 latent variable로 취급하여 학습을 진행한다.
참고: https://ploradoaa.tistory.com/107

1. Models

RAG-Sequence Model

RAG-Sequence 모델은 전체 sequence를 생성하기 위해 동일한 검색된 문서를 사용하는 방법이다. 수식을 살펴보면, $x$가 주어졌을 때 document $z$를 생성할 확률을 계산하고, 이 하나의 $z$를 가지고 모든 target sequence $y$를 생성할 확률을 $z$에 대해 marginalize해주는 과정임을 볼 수 있다. 이 때 상위 K document에 대해서만 margnalize를 진행한다.

RAG-Token Model

RAG-Token 모델은 각 target token을 생성할 때 서로 다른 latent document를 가져올 수 있다. Token 단위로 생성된다고 볼 수 있다. Generator는 답변을 생성할 때 여러 문서에서 content를 선택할 수 있다. 상위 K document는 retriever를 통해 검색되고, generator는 각 document에 대한 다음 output token의 distribution 생성, marginalize를 진행하여 token을 생성한다. 이런 방식으로 다음 output token에 반복적으로 적용한다.

2. Retriever: DPR

Retriever $p_\eta(z∣x)$는 bi-encoder architecture인 DPR을 기반으로 한다.

$d(z)$: $\text{BERT}_{\text{BASE}}$구조 document encoder에서 생성된 dense representation
$q(x)$: $\text{BERT}_{\text{BASE}}$구조 query encoder에서 생성된 query representation

위의 식과 같이 $p_\eta(z∣x)$는 $d(z)$와 $q(x)$의 내적에 의해 나타낼 수 있고, 내적 값이 높을 수록 선택될 확률이 높다. Prior probability $p_\eta(z∣x)$에서 $k$개의 document를 계산하는 문제는 Maximum Inner Product Search (MIPS) 문제로, sub-linear 시간안에 해결할 수 있다.

DPR의 pre-trained bi-encoder를 사용하여 retriever를 초기화하고 document index를 구축했다. 이 document index를 non-parametric 메모리라고 한다.

3. Generator: BART

Generator $p_\theta(y_i ∣ x, z, y_{1:i-1})$ 로는 어떤 encoder-decoder도 사용가능하다. 본 논문에서는 BART-large를 사용했으며, input $x$와 검색된 $z$를 결합해서 BART 입력으로 넣어주기 위해 단순 concat을 사용했다.

4. Training

앞서 언급했듯 Retriever와 Generator는 동시에 학습되고, 입력 $x$ ↔ target $y$에 대해 supervised 학습을 진행하지만 document $z$를 구하는 과정에는 따로 직접적인 정답이 주어지지 않는다.
Loss function: target에 대한 negative marginal log-likelihood: $\sum_j - \log p(y_j ∣ x_j)$
REALM📄에서는 document encoder도 학습 진행하였는데, 이 과정은 document index를 매번 update해야하므로 cost가 많이 든다. 저자는 document encoder 학습이 성능에 큰 영향을 미치지 않는다고 판단, document encoder는 freeze하고 query encoder $\text{BERT}_q$ and the BART generator만 fine-tuning을 진행했다.

5. Decoding

Test time에 RAG-Sequence와 RAG-Token은 output decoding, 즉 $\arg \max_y p(y|x)$를 구하는 데도 서로 다른 방법이 필요하다.

RAG-Token

RAG-Token 모델은 token별로 확률을 구하므로 아래 수식과 같이 일반적인 autoregressive seq2seq generator와 동일하게 나타낼 수 있다. 따라서 일반적인 beam decoder를 사용하여 decoding할 수 있다.

Beam Search
기본적으로 Seq2Seq 모델에서 디코딩을 진행할 때, 각 시점에서의 확률이 가장 높은 것을 채택하는 Greedy Decoding 방식을 주로 사용한다. 하지만 이 방법은 특정 시점의 확률분포에서 미묘한 차이로 1등을 차지하더라도 2등은 고려되지 않는다. 시간 복잡도 측면에서는 유리하지만, 이전 예측이 다음 예측에 중요한 영향을 미치는 autoregressive decoding에서 이러한 방식은 최종 예측의 정확도를 떨어뜨릴 수 있다.

Beam Search는 이러한 문제점을 개선하기 위해 상위 $K$개(=beam의 개수)의 경우의 수를 고려하는 방법이다. 모든 경우의 수를 고려하기엔 시간 복잡도 측면에서 불가능하기 때문에 상위 K를 선정한다.

(1) 시작 token을 바탕으로 나온 예측 확률 중 가장 높은 $K$개 token을 고르고, (2) 이 K개를 갈래로 또 각각 $K$개를 예측한다. (3)이 때, 전체 $K^2$개의 갈래 중 누적 확률을 기준으로 상위 $K$를 뽑는다.

고려해야할 점은, 고려하는 모든 확률은 누적 확률이므로 특정 계층의 하위 노드들이 서로 같은 확률을 가지더라도, 어떤 빔에서 뻗어나왔냐에 따라 누적 확률은 달라지게 된다는 점이다. 이와 같은 방식으로 $K$개의 beam을 유지해가며 종료 token을 만날 때 까지 반복적으로 decoding을 진행한다.
참조: https://sooftware.io/beamsearch/

RAG-Sequence

반면에 RAG-Sequence 모델은 token별 likelihood로 분리되지 않고, 하나의 document에 대해 sequence 전체를 생성하므로 token을 하나씩 생성하며 확률을 고려하는 beam search 방법을 적용할 수 없다.

따라서 논문에서는 각 document $z$에 대해 beam search를 진행하는 방법을 사용했다. 방법은 다음과 같다.

먼저 top K개의 document $z$를 검색하고, 각 document에 대해 beam search를 적용하여 top K sequence를 뽑는다. 그럼 각각의 output에 대한 확률은 위 그림의 오른쪽과 같이 나타낼 수 있다. 이를 document $z$에 대해 marginalize를 진행하여 $p(y∣x)$를 구해야하는데, 이 때 문제가 발생한다.

예를 들어 output $y1$을 예시로 살펴보자. 해당 식에서 $p(y∣x,z2)$는 beam search 과정에서 구할 수 없는 값이다. $z2$로부터는 $y2, y3, y4$만 생성되었기 때문에 $y1$에 대한 값은 구할 수 없다.

이러한 문제를 해결하기 위해, 두 가지 방법을 적용했다.

Thorough Decoding: beam search에서 발견되지 않은 값들을 위한 추가적인 forward pass(해당 확률을 추가적으로 계산)를 통해 해당 확률을 구함. 이 방법은 계산량이 많아 효율적이지 않음.
Fast Decoding: 발견되지 않은 값들은 0으로 처리. 효율적인 계산이 가능.

Experiments

1. Open-domain Question Answering

Open-domain question answering(QA)은 real-world application이자 knowledge-intensive task를 위한 일반적인 task이다. 논문에서는 질문과 답변을 input-output text pairs (x, y)로 설정하고 답변의 negative log-likelihood를 최소화하여 RAG를 학습했다. 여기서 답변은 검색된 document에서 추출된 범위이며 non-parametric knowledge에 의존한다.

또한 검색을 활용하지 않고 parametric knowledge에만 의존하는 “Closed-Book QA” 접근방식과 비교했다.

데이터셋은 총 4가지를 사용했다: Natural Questions(NQ), TriviaQA(TQA), WebQuestions(WQ), CuratedTrec(CT)

결과는 위와 같다. 4개의 Open-domain QA task에서 RAG가 모두 SOTA를 달성한 것을 확인할 수 있다. RAG는 “closed-book”(parametric only) 접근 방식의 생성 유연성과 “open-book” 검색 기반 접근 방식의 성능을 결합한다.

2. Abstractive Question Answering

RAG 모델은 간단한 extractive QA를 넘어 free-form 추상적 텍스트 생성으로 질문에 답할 수 있다. knowledge-intensive 환경에서 RAG의 자연어 생성(NLG)을 테스트하기 위해 MSMARCO NLG task v2.1 데이터셋을 사용했다. 이 작업은 질문, 각 질문에 대해 검색 엔진에서 검색한 10개의 gold passage, 검색된 구절로부터 주석이 생성된 full sentence 답변으로 구성된다. 제공된 구절을 사용하지 않고 질문과 답변만 사용하여 MSMARCO를open-domain abstractive QA task으로 처리한다.

표 2에서 볼 수 있듯이, RAG-Sequence는 BART보다 높은 성능을 보인다. RAG는 SOTA 성능에 근접한 결과를 보이는데, 여기서 (i) SOTA 모델이 reference 답변을 생성하는 데 필요한 특정 정보가 있는 gold passage에 액세스한다는 것, (ii) 많은 질문이 gold passage 없이는 답변할 수 없다는 점, (iii) 모든 질문에 Wikipedia만으로 답변할 수 있는 것은 아니라는 점을 감안할 때 의미가 있다.

또한 표 3은 RAG 모델에서 생성된 일부 답변이다. 정성적으로, RAG 모델은 BART보다 hallucination이 적고 사실적으로 정확한 텍스트를 더 자주 생성하는 것을 볼 수 있다.

3. Jeopardy Question Generation

Non-QA 환경에서 RAG의 생성 능력을 평가하기 위해 open-domain question generation에 대해 실험을 진행했다. 일반적으로 짧고 간단한 질문으로 구성된 표준 open-domain QA task의 질문을 사용하는 대신 Jeopardy 질문 생성이라는 더 까다로운 task를 사용했다. Jeopardy는 특정 entity에 대한 사실에서 entity를 추측하는 형식이다. 예를 들어, “월드컵”은 “1986년 멕시코가 이 국제 스포츠 대회를 두 번 개최한 최초의 국가로 기록되었습니다.”라는 질문에 대한 답이다.

평가 metric으로는 Q-BLEU-1(question generation에 대한 인간 판단과의 상관 관계가 더 높은 BLEU의 변형)을 사용했다. 또한 generation 사실성 평가, 구체성을 평가에 대해 두 가지 human evaluation을 수행했다.

표 2를 통해 RAG-Token이 Jeopardy 생성에서 RAG-Sequence보다 더 나은 성과를 보이며, 두 모델 모두 Q-BLEU-1에서 BART보다 더 나은 성과를 보이는 것을 확인할 수 있다.

표 4는 BART와 RAG-Token의 human evaluation 결과를 나타낸다. SOTA 생성 모델보다 RAG가 효과적임을 보여주고 있다.

Jeopardy 문제에는 두 개의 별도 정보가 포함되어 있는 경우가 있다. RAG-Token은 여러 문서의 콘텐츠를 결합한 응답을 생성할 수 있기 때문에 Jeopardy 문제에서 가장 좋은 성과를 보인다. 그림 2는 그 예를 보여준다. “Sun”을 생성할 때, document 2의 posterior가 높다. 마찬가지로 “A Farewell to Arms”을 생성할 때 document 1의 posterior가 더 높다. 흥미롭게도 각 book의 첫 번째 token이 생성된 후 document posterior가 평평해진다. 이 결과는 generator가 특정 document에 의존하지 않고도 title을 완성할 수 있음을 나타낸다. 즉 parametric 메모리와 non-parametric 메모리가 함께 작동하는 방식을 보여준다.

AI, Paper Review

This post is licensed under CC BY 4.0 by the author.