REALM: Retrieval-Augmented Language Model Pre-Training

논문 개요

대규모 언어 모델은 사전학습을 통해 방대한 사실 지식(factual knowledge)을 파라미터에 암묵적으로 저장한다. 그러나 이 접근법에는 두 가지 근본적 한계가 있다. 첫째, 지식이 파라미터에 분산 저장되므로 특정 사실이 어디에 어떻게 저장되어 있는지 해석하기 어렵다. 둘째, 세상의 지식은 변하는데, 모델의 파라미터는 학습 이후 업데이트하지 않으면 구식 정보를 담게 된다.

Guu et al.(2020)이 ICML 2020에서 발표한 REALM(Retrieval-Augmented Language Model Pre-training)은 이 문제를 사전학습 단계에서 해결하는 새로운 패러다임을 제시한다. 핵심 아이디어는 사전학습 중에 LM이 예측을 내릴 때마다 외부 문서 코퍼스에서 관련 문서를 검색하여 맥락으로 활용하는 것이다. 이를 통해 지식이 파라미터가 아닌 외부 저장소에 명시적으로 유지되며, 저장소를 업데이트하면 모델 재학습 없이도 지식을 갱신할 수 있다.

---

핵심 기여

1. 사전학습 단계 검색 통합: 파인튜닝이나 추론 시점이 아닌 사전학습 단계에서 검색을 통합하여, 모델이 처음부터 검색을 활용하도록 학습한다.
2. 엔드투엔드 공동 학습: 지식 검색기와 언어 모델이 단일 역전파로 함께 최적화된다. 검색기가 더 유용한 문서를 찾도록 자동으로 학습된다.
3. MIPS 비동기 인덱스 업데이트: 학습 중 MIPS 인덱스를 비동기적으로 주기적 갱신하여, 매 스텝마다 인덱스를 재구축하는 계산 비용을 회피한다.
4. 해석 가능한 지식 저장: 모델이 어떤 문서를 참조하여 예측을 내렸는지 추적할 수 있어, 지식 출처가 명확해진다.
5. Open-QA SOTA: Open-domain QA 벤치마크(NQ, WebQ, TriviaQA)에서 당시 최고 성능을 달성하였다.

---

방법론 상세

모델 구조

REALM은 두 가지 핵심 구성 요소로 이루어진다:

지식 검색기 (Knowledge Retriever) $p(z|x)$:

• 입력 $x$에 대해 문서 코퍼스 $\mathcal{Z}$에서 관련 문서 $z$를 검색
• BERT 기반의 양방향 인코더(bi-encoder) 구조: ${\text{Embed}}_{\text{input}}(x)$와 ${\text{Embed}}_{\text{doc}}(z)$를 각각 계산

$p(z|x)=\frac{\exp (f(x,z))}{{\sum }_{z^{\prime }}\exp (f(x,z^{\prime }))},f(x,z)={\text{Embed}}_{\text{input}}(x{)}^T{\text{Embed}}_{\text{doc}}(z)$

Knowledge-Augmented Encoder $p(y|z,x)$:

• 검색된 문서 $z$와 입력 $x$를 결합하여 출력 $y$를 예측
• BERT를 사용하여 $[\text{CLS}]x[\text{SEP}]z[\text{SEP}]$를 처리

사전학습 목표

REALM의 사전학습은 마스킹 언어 모델링(MLM)을 검색과 결합하여 수행된다:

${\mathcal{L}}_{\text{REALM}}=-\log p(y|x)=-\log {\sum }_{z\in {\mathcal{Z}}_{\text{top-k}}}p(z|x)\cdot p(y|z,x)$

여기서 $y$는 마스킹된 토큰의 원래 값, $x$는 마스킹이 적용된 입력 문장이다. 문서가 도움이 될 경우 $p(z|x)\cdot p(y|z,x)$가 크므로 검색기가 해당 문서를 더 높이 랭킹하도록 학습된다.

MIPS와 비동기 인덱스 업데이트

코퍼스에서 관련 문서를 효율적으로 검색하기 위해 최대 내적 탐색(MIPS: Maximum Inner Product Search)을 사용한다:

$z^{\ast }=\arg {\max }_{z\in \mathcal{Z}}f(x,z)=\arg {\max }_{z\in \mathcal{Z}}{\text{Embed}}_{\text{input}}(x{)}^T{\text{Embed}}_{\text{doc}}(z)$

학습 중에 ${\text{Embed}}_{\text{doc}}$의 파라미터가 변하면 인덱스도 갱신해야 한다. 그러나 매 학습 스텝마다 전체 코퍼스의 임베딩을 재계산하는 것은 불가능하다.

비동기 업데이트 전략:

1. 학습 초기: 인덱스를 무작위 초기화로 구성
2. 매 $T$ 스텝(예: $T=500$)마다 비동기 프로세스가 전체 코퍼스 임베딩을 재계산하여 인덱스 갱신
3. 인덱스가 약간 stale(구식)하더라도 학습에 큰 영향이 없음을 실험적으로 확인

이를 통해 MIPS 인덱스 갱신 비용은 학습 FLOPs의 약 $1/T$ 비율이 된다.

사전학습 데이터 설계

입력 코퍼스 $\mathcal{Z}$: Wikipedia의 모든 문서를 사용 (약 3억 단어, 10만 개 청크로 분할)

마스킹 전략: 일반 토큰 마스킹 대신 개체명(entity) 마스킹을 사용한다. 문장에서 개체명(인물, 장소, 조직 등)을 마스킹하면 모델이 사실적 지식(factual knowledge)을 검색하도록 유도된다. 예:

“[MASK]는 1990년에 노벨 물리학상을 수상하였다.” → 검색기가 해당 수상자에 관한 Wikipedia 문서를 찾아야 함

파인튜닝 (Open-QA)

사전학습된 REALM을 Open-QA에 파인튜닝할 때:

$p(y|x)={\sum }_{z}p(z|x)\cdot p(y|z,x)$

여기서 $p(y|z,x)$는 BERT 판독기(reader)가 문서 $z$에서 답변 $y$를 추출하는 확률이다. 상위 $k$개(보통 $k=5$) 문서에 대해 합산한다.

---

실험 결과

Open-QA 벤치마크 성능 (Exact Match)

모델	NaturalQuestions	WebQuestions	TriviaQA
T5-11B (파라미터에 지식 저장)	34.5	-	-
DPR (DPR+BERT 판독기)	41.5	41.7	57.9
REALM	40.4	40.7	-
REALM (NQ fine-tuned)	44.5	43.5	-

당시 기준으로 Open-QA SOTA를 달성하였으며, 특히 파라미터에 지식을 저장하는 방식(T5-11B)보다 훨씬 적은 파라미터(300M)로 더 높은 성능을 보였다.

사전학습의 중요성

검색 통합을 파인튜닝 단계에만 적용한 경우(REALM w/o pre-train retrieval)와 비교하면 NQ에서 약 3%p 성능 차이가 있었다. 사전학습 단계에서의 검색 통합이 핵심임을 보여준다.

개체명 마스킹의 효과

일반 토큰 마스킹 대신 개체명 마스킹을 사용할 때 NQ 성능이 약 2%p 향상되었다.

---

의의 및 한계

의의

• RAG의 선구자: REALM은 이후 RAG(Retrieval-Augmented Generation, Lewis et al., 2020), Atlas(Izacard et al., 2022) 등 검색 증강 생성 연구의 직접적인 출발점이 되었다.
• 명시적 지식 저장: 지식이 파라미터가 아닌 외부 저장소에 있으므로 업데이트, 수정, 디버깅이 용이하다.
• 해석 가능성: 검색된 문서를 통해 모델의 예측 근거를 추적할 수 있다.
• 파라미터 효율: 대형 모델 없이도 방대한 지식을 활용할 수 있다.

한계

• 학습 복잡성: 검색기와 언어 모델을 공동 학습하는 것이 표준 LM 학습보다 훨씬 복잡하다. 특히 비동기 인덱스 업데이트 구현이 어렵다.
• 검색 지연: 추론 시 매 예측마다 MIPS 검색이 필요하므로, 순수 파라메트릭 모델보다 추론 속도가 느리다.
• 검색 실패 전파: 검색기가 관련 없는 문서를 반환하면 최종 예측도 오류가 된다. 검색 오류가 증폭되는 구조이다.
• 코퍼스 크기 제약: 검색 코퍼스가 클수록 좋지만, MIPS 인덱스 메모리와 검색 지연이 함께 증가한다.
• 생성 과제 한계: REALM은 추출형 QA에 최적화되어 있으며, 자유 형식 생성(open-ended generation)에는 직접 적용이 어렵다. 이 한계는 이후 RAG와 Atlas에서 개선된다.