원제: Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

잡소리: 원래 BM-25한 김에 SPLADE를 할려고 했는데, 아무리봐도 이게 먼전거같아서 RAG를 우선 읽어본다.
굳이 순서를 따지자면 연도순대로 BM25 - DPR - RAG - SPLADE나 아니면 DPR - RAG - BM25 - SPLADE로 읽도록 하자

1. Introduction

Pre-training Neural Language Model 자체도 외부 메모리에 대한 접근 없이 지식 베이스로서 동작이 가능하다. 하지만 메모리에 대한 확장이나 수정이 어렵고, 직관적인 설명도 제공할 수 없고, 환각(hallucination)을 생성하기도 한다.

이를 해결하기 위해 parametric memory와 non-parametric memory (retreival-based)를 결합한 하이브리드 모델이 나온다.(REALM, ORQA) 하지만 부분 Open domain에서 추출해오는 QA형 과제에만 적용되기에, 해당 모델에서는 Seq2Seq 모델로 확장시킨다.
즉, parametric memory로 pre-training된 Seq2Seq transformer를 사용하고, non-parametric memory로 위키피디아 Dense vector index를 활용하여 검색하는 DPR를 활용해 두 요소를 하나의 확률 모델로 결합하여 end-to-end로 학습한다.

기존이 연구들과 다른 점은 parametric memory와 non-parametric memory 모두 pre-training되어있기에 추가적인 학습 없이도 지식에 접근할 수 있는 능력이 존재한다.

2. Method

입력 시퀀스 $x$, 사용할 텍스트 문서 $z$, 목표 시퀀스(생성) $y$
이때 DPR논문에서처럼
1) Retriever: Question $x$에 대해 text passage에 대한 분포를 변환하는 매개변수 $\eta$를 갖는 $p_{\eta}(z|x)$
2) Generator: 이전 i-1개 토큰 $y_{1:i-1}$, 입력 $x$, 검색된 passage $z$에 따라 현재 토큰을 생성하는 매개변수 $\theta$를 가지는 $p_{\theta}(y_i | x, z, y_{1:i-1})$

end-to-end를 위해 Retriever을 통과한 문서를 잠재 변수(latent variable) 취급하고, 주변화하여 생성된 텍스트에 대한 분포를 generate하는 두 가지 모델이 존재한다.

Model

RAG-Sequence

동일한 문서를 사용하여 전체 목표 시퀀스 예측하는 것이 목표이다.
검색된 문서를 하나의 잠재 변수로 취급하고, 이를 marginalize하여 Seq2Seq 확률 $p(y|x)$를 구한다.
$p_{RAG\_Sequence}(y|x) \approx \displaystyle\sum_{z \in top-k(p(\cdot|x))}{p_\eta (z|x)p_\theta(y|x, z)} = \displaystyle\sum_{z \in top-k(p(\cdot| x))}p_\eta (z|x) \prod_{i}^{N} p_\theta(y_i|x,z,y_{1:i-1})$

식이 길고 지루하고 현학적이다. 장난이고 조금 더 자세히 설명해보면 $p_\eta(z|x)$는 질문 x에 대해 문서 z가 관련 있을 확률, $p_\theta(y|x,z)$는 x, z를 참고해 y라는 답변이 나올 확률이다. 즉, (해당 문서를 썼을 확률) $\times$ (이 문서를 통해 그런 답이 나올 가능성)으로 해석하면 된다. 가장 오른쪽 식은
$p_\theta(y|x, z) = \displaystyle\prod_{i}^{N}p_\theta(y_i|x,z,y_{1:i-1})$ 즉, 전체 문장이 나올 확률은 각 토큰이 나올 확률의 곱과 같다는 점만 생각해보면 된다.

RAG-Token

각 목표 토큰마다 서로 다른 잠재 문서를 선택해서 이에 대해 marginalize한다. 즉, retriever이 가져온 top-K 문서에서 선택적으로 활용하여 각 문서에 대해 다음 출력 토큰의 분포를 계산하고 marginalize한다.
$P_{RAG\_Token}(y|x) \approx \displaystyle\prod_{1}^N \displaystyle\sum_{z\in top-k(p(p\cdot|x))} p_\eta(z|x)p_\theta(y_i|x,z,,y_{1:i-1})$

DPR

자세한건 이전 글을 참고하자
$p_\eta(z|x) \propto exp(d(z)^{\top}q(x))$, 이때 $d(z) = BERT_d(z)$, $q(x)=BERT_q(x)$

즉, d(z)는 문서의 dense, q(x)는 Query 표현 이다.
$top-k(p_\eta(\cdot |x))$를 계산하는 것은 MIPS문제로, 이는 sub-linear time동안 근사적으로 계산이 가능하다. (대충 그럴듯한 top-K를 찾는다고 보면 된다)

해당 문서 인덱스를 non-parametric memory라고 부른다.

Generator: BART

Generator은 다른 encoder-decoder모델로도 가능하지만, 해당 논문에서는 pre-training된 seq2seq인 BART-large를 사용한다.

출력 생성시에 입력 x와 Retriever를 통과한 문서 z를 단순하게 concat한다.
BART generator의 매개변수 $\theta$는 parametric memory라고 부른다.

Training

Retriever와 Generator를 공동으로 학습하며, 입력-출력 쌍에 대해 negative marginal log-likelihood를 Adam을 사용하는 확률적 경사 하강법으로 최소화한다.
이때 학습 중에 $BERT_d$를 업데이트하는 것은 비용은 크지만 필수적이지 않기에 문서 encoder와 문서 인덱스를 고정하고 $BERT_q$와 BART Generator만 fine-tuning한다.

Decoding(토큰을 선택해서 문장 만들기)

테스트 시점에서 $argmax_y p(y|x)$를 근사하기 위한두 모델 간 차이가 약간 존재한다.
우선 Standard beam decoder을 위해선 다음 토큰의 확률이 이전 토큰들에만 의존하고, 토큰마다 하나의 명확한 확률 분포가 필요하다.

beam search: decoding 단계에서 확률이 제일 큰 토큰만 고르면 문장이 망할 수 있기에 여러 가능성을 고려하는 전략

RAG-Token

$P_{RAG\_Token}(y|x) \approx \displaystyle\prod_{1}^N \displaystyle\sum_{z\in top-k(p(p\cdot|x))} p_\eta(z|x)p_\theta(y_i|x,z,,y_{1:i-1})$

해당 식에서
$p'_\theta(y_i|x,y_{1,i-1}) = \displaystyle\sum_{z \in top-k(p(\cdot|x))}p_\eta(z|x)p_\theta(y_i|x,z,y_{1:i-1})$로 사실상 표준 Seq2Seq generator로 볼 수 있기에 Standard beam decoder을 그대로 사용 가능하다.

RAG-Sequence

$p(y|x)$가 토큰 단위의 likelihood로 분해되지 않는다. 따라서 문서 z에 대해 개별적으로 beam search를 수행한다. 이때, 가설 y에 대해
$p_\theta(y_i | x,z,y_{1:i-1})$를 사용해 점수화된다. 이 과정에서 가설 집합 Y를 얻는데, 일부는 beam 결과에 나타나지 안흔ㄴ다.

가설 y의 확률을 추정하기 위해 해당 가설이 beam에 등장하지 않은 문서 z 각각에 대해 추가적인 forward pass를 수행하고 Generator 확률에 $p_\eta (z|x)$를 곱한 후 모든 문서에 대해 확률을 합산하여 margin 확률을 계산한다.

만약 출력 시퀀스가 길어서 |Y|가 너무 커진다면 beam search 과정에서 생성되지 않은 y에 대해 $p_\theta(y|x, z_i) \approx 0$이라고 근사해버린다.

3. Experiments & Result

DPR때처럼 위키피디아를 100단어 chunk기준 passage로 분할해서 2100만개 문서를 만들고, FAISS를 이용한 HNSW 기반의 MIPS를 구축하여 검색을 수행한다.

Open-domain Question

QA의 질문-답변을 입력-출력 쌍으로 취급하고 negative log-likelihood를 최소화하도록 RAG를 학습한다. 널리 사용되는 NQ, TQA, WQ, CT 데이터셋을 이용한다.

결과에서 크게 중요한 것은
1) 특수한 salient span masking pre-training없이도 강력한 성능을 보인다.

salient span masking: 아무 span이나 가리는 것이 아닌 핵심 정보만 가리는 방법

2) Re-ranker나 extractive reader 없이도 경쟁력있는 성능을 보이며 둘이 필수가 아님을 보여준다.

Abstractive Question Answering

Knowledge-intensive setting에서 자연어 생성 능력을 평가하기 위해 MSMARCO NLG v2.1을 이용했다. 이 중 특정 지역의 날씨와 같이 정답 문서에 접근하지 않으면 일치하는 답변을 할 수 없는 질문도 포함되어있다.

결과를 요약해보면 기존의 BART에 비해서 환각이 적고 사실적으로 올바른 텍스트를 생성하는 비율이 높다

Jeopardy Question

QA가 아닌 설정에서의 RAG의 생성 능력을 평가한다. Jeopardy는 어떤 개체에 대한 사실이 주어지고, 사실에 해당하는 답을 맞히는 형식이다. 정확하고 사실적인 진술이 요구되기에 answer entities를 조건으로 질문을 생성하는 것은 매우 어렵다고 한다.

결과를 요약해보면
1) Jeopardy 질문은 서로 다른 정보 조각을 포함하는 경우가 있는데, RAG-Token은 여러 문서에서 가져온 내용을 결합할 수 있기에 더 나은 성능을 보였다.

2) 제목의 일정 부분까지 보고난 후 모든 문서의 확률이 유사해지는 현상이 나온다. 즉, 검색 문서 없이도 다음 토큰을 예측한다. 이는 특정 지식이 parameter에 저장된다는 것을 알 수 있다. 즉, non-parametric memory가 생성을 유도하고 parametric-memory가 완성하는 방식으로 두 메모리가 협력함을 보인다.

Fact Verfication

FEVER task는 자연어 주장에 대해 위키피디아에 의해 지지/반박/정보 부족인지를 분류하는 태스크이다. 즉, 생성이 아닌 분류가 가능한지를 판단한다.

결과는 gold evidence(정답이 있는 파일)을 잘 가져오고, 설령 gold evidence가 없더라도 정확도가 크게 흔들리지 않았다.

추가 결과

1) BART보다 더 다양한 출력을 생성한다(RAG-Sequence > RAG-Token)

2) 제거 실험을 진행한 결과 학습된 검색이 성능을 향상시킴을 증명했고, FEVER에선 BM25와 같은 inverted-index를 retriever로 사용하는게 더 좋았지만, 그 외의 다른 태스크에선 differentiable retrievall이 더 효과적이었다.

3) Non-parametric memory의 교체만으로 모델의 지식을 업데이트 할 수 있었다.

4) 학습시 k에 따른 성능 차이는 크지 않았고, test-time에는 상황에 따라 다른 k값을사용할 수 있다(Retriever가 확률분포이기에).
RAG-Sequence의 경우 k가 늘어날수록 성능이 단조증가했고(50근처까지), RAG-Token은 k가 일정 이상 증가하면 오히려 noise가 심해지며 성능이 감소했다.

Conclusion

생각보다 Titans때 처럼 엄청 식이 복잡하거나 그런 것은 거의없었다. 뭔가 이해가 될랑말랑해서 내가 쓴 정리본좀 다시 몇번 읽어봐야 될 것 같다. 그리고 아마 다음 SPLADE가 단점을 또 보완한 부분이 있을거라 이 방식의 어떤 단점이 있었을지도 고민해봐야겠다.