Background
소속: Google Research
출판: PMLR 2020
Abstract & 1. Introduction
LLM은 뛰어난 성능을 지니지만 모든 지식들은 파라미터에 저장되고, 더 많은 지식을 다루기 위해서는 점점 네트워크를 키워야한다. latent knowledge retriever를 활용해 모듈화하고 해석 가능하도록 지식을 습득하기 위한 새로운 프레임워크 Retrieval-Augmented Language Model(REALM)을 제안한다. 이 모델은 추론 과정 중에 어떤 지식을 사용하고 검색할 지 모델이 결정하게 함으로써 지식의 역할을 명시적으로 노출한다.
모델을 학습하기 위해 마스킹된 문장의 마스크를 예측하는 unsupervised 학습 방법을 채택한다. retriver이 마스킹에 도움이 되는 문서를 찾아내면 보상을 주고 아니라면 불이익을 준다. Figure 1과 같이 마스킹된 언어로부터 문서를 검색하고 마스크를 예측하고, 적절한 마스크를 예측했는지에 따라 적절한 문서 검색이 이뤄졌는지 평가하는 과정을 end-to-end로 훈련한다. 한편 매번 retriver이 수백만개의 문서를 검색하지 않고, 이미 저장된 문서를 기반으로 검색을 수행하며 비동기적으로 업데이트한다. 최고의 문서는 Maximum Inner Product Search(MIPS)를 계산하여 구한다.
3. Approach
3.1. REALM’s generative process
pre-training과정과 fine-tunging과정 모두 입력 $x$가 주어지면 출력 $y$을 생성할 확률 분포 $p(y \mid x)$를 예측한다.
- pre-training: $x$는 일부 토큰이 마스킹된 문장 corpus $X$에 포함된 문장이고 $y$는 마스킹된 토큰.
- fine-tuning: $x$는 질문이고 $y$는 그 답변
REALM은 $p(y \mid x)$를 두 단계로 분해할 수 있다.
- retrieve:입력 $x$가 주어지면, knowledge corpus $Z$로부터 유용한 문서 $z$를 검색할 확률 분포 $p( z \mid x)$를 예측한다.
- predict: 그리고 입력 $x$와 문서 $z$가 주어졌을 때 출력 문장 $y$가 생성될 확률 분포 $p(y \mid z, x)$를 예측한다.
$y$를 생성할 확률을 얻기 위해 $z$를 latent variable로 다루고 가능한 모든 문서 $z$에 대해 적분하여 주변화(marginalization)한다.
$$ p(y \mid x) = \sum_{z \in Z} p(y \mid z, x) p(z \mid x) $$
Marginalize (주변화)는 여러 확률 변수의 결합 확률 분포가 주어졌을 때, 특정 변수의 확률 분포 (주변 확률 분포, marginal probability distribution)를 구하기 위해 다른 변수들에 대해 합하거나 (이산 변수의 경우) 적분하는 과정
주변 확률 분포 (Marginal Probability Distribution): 특정 변수 하나만의 확률 분포를 의미하며, 다른 변수들의 값에 상관없이 해당 변수가 특정 값을 가질 확률
Marginalize의 의미와 중요성:
- 관심 변수에 집중: 여러 변수가 관련된 복잡한 확률 분포에서 특정 변수의 확률적 특성만 파악하고자 할 때 유용합니다. 다른 변수들의 영향을 제거하고 관심 변수에 대한 순수한 확률 분포를 얻을 수 있습니다.
- 정보 축약: 고차원의 결합 분포를 저차원의 주변 분포로 축약하여 분석을 단순화할 수 있습니다.
3.2. Model architecture
Knowledge Retriever
The retriever is defined using a dense inner product model:
$$ p(z \mid x) = \frac{\exp{f (x, z)}}{\sum_{z'} \exp{f (x, z^{'})}}$$
$$f(x, z) = \text{Embed}_{\text{input}}(x)^\top \text{Embed}_\text{doc}(z)$$
$$ \text{Embed}_{\text{input}}(x) = \text{BERT}_{CLS}(x), text{Embed}_\text{doc}(z) = \text{BERT}_{CLS}(z_{title}, z_{body})$$
- $f(x, z)$: relevance score. 임베딩된 두 벡터의 내적.
- $p(z \mid x)$: retrieval distribution. 모든 relevance scores에 대한 softmax 결과
- $z_{title}$ is the document’s title and $z_{body}$ is its body.
Knowledge-Augmented Encoder
입력 $x$와 retrieved document $z$가 주어졌을때 $y$를 예측할 확률 분포 $p(y \mid z, x)$를 정의한다.
먼저 $x$와 $z$를 한 시퀀스로 묶어 Transformer에 입력한다. 이를 통해 $y$를 예측하기 전에 $x$와 $z$의 풍부한 cross-attention을 수행할 수 있다.
이 과정에서는 pre-training과 fine-tuning과정이 조금 다르다. pretrining에서는 $x$에서 마스킹된 토큰을 예측하고 fine-tuning에서는 답변을 생성한다.
- pretraining: $$p(y \mid z, x) = \prod_{j=1}^{J_x} p(y_j \mid z, x)$$ $$ p(y_j \mid z, x) \propto \exp(w_j^\top \text{BERT}_{\text{MASK}(j)} (\text{join}_{\text{BERT}}(x, z_{body})) $$
- fine-tuning: $$ p (y \mid z, y) = \sum_{s \in S(z, y)} \exp(\text{MLP} ([h_{\text{START}(s)}; h_{\text{END}}(s)])$$ $$ h_\text{START}(s) = \text{BERT}_{\text{START}(s)}(\text{join}_\text{BERT}(x, z_{body}) $$ $$ h_{\text{END}(s)} = \text{BERT}_{\text{END}(s)}(\text{join}_{\text{BERT}}(x, z_{body})$$
3.3. Training
pre-training과 fine-tuning 모두 x가 주어졌을 때 정답 y가 생성될 확률$ \text{log} p(y \mid x)$을 최대화하는 방향으로 진행된다. knowledge corpus $\mathcal{Z}$에 포함된 모든 documents $z$에 대한 확률을 계산하는 것은 어려운 문제이기 때문에 확률이 높다고 예측되는 상위 k개의 document에 대해서만 $p(z \mid x) $에 대한 최적화를 진행한다. 이때 상위 k개는 relevance score $f(x, z) = \text{Embed}_{\text{input}}(x)^\top \text{Embed}_{\text{doc}}(z)$를 내림차순으로 정렬하여 결정한다. 이러한 방식을 적용하기 위해 모든 $\text{Embed}_{\text{doc}}(z)$에 대한 index가 미리 계산되어있어야한다. 업데이트된 모델에 맞게 index도 정확한 정보를 가져야하지만, 매번 업데이트하기에는 비용이 많이 들고 어차피 상위 k개의 문서만을 사용하기 때문에 비동기적으로 몇 백번의 training step마다 index를 업데이트한다.
What does the retriever learn?
$f(x, z)$를 기반으로 $x$와 $z$가 얼마나 연관이 있는지 스코어링을 할 수 있다. 이러한 스코어링의 정확도를 높이기 위해 다음의 수식을 통해 파라미터를 업데이트한다.
$$ \nabla \log p(y \mid x) = \sum_{z \in \mathcal{Z}} r(z) \nabla f(x, z) $$
$$ r(z) = [\frac{ p(y \mid z, x)}{p(y \mid x)}-1]p(z \mid x)$$
- $ \nabla f(x, z)$: 각 문서 $z$가 주어진 입력 $x$에 대해 어떤 영향을 주는지를 측정하는 그래디언트. "그 문서를 통해 y가 더 잘 생성되게 하려면 모델의 파라미터를 어떻게 바꿔야 할까?"라는 방향을 나타냄
- $ \frac{ p(y \mid z, x)}{p(y \mid x)} -1 $:$x$만으로 $y$를 예측할 때보다 $x$와 $z$가 주어졌을 때 $y$를 예측하는 게 얼마나 더 확률이 높은가에 대한 계산.
- $ [\frac{p(y \mid z,x)}{p(y \mid x)}-1] p(z \mid x)$: $p(z \mid x)$는 문서가 선택될 확률. 선택되었을 때 $y$를 예측할 확률이 높은데 선택될 확률도 높으면 곱해줌으로써 크게 반영하고 아니면 적게 반영할 수 있도록 가중치를 주는 역할.
- $r(z)$: 기여 정도 $\times$ 선택될 확률
3.4. Injecting inductive biases into pre-training
실제 학습하게 될 때 살펴봐야할 내용들
4. Experiments
4.1. Open-QA Benchmarks
- NaturalQuestions-Open: 구글 쿼리와 답변. 본 논문에서는 최대 다섯개의 토근으로 이뤄진 "short answer type"만 분류함. 데이터셋에서 추천할 만한 Wikipedia 문서도 제공하지만, 우리 모델에서는 해당 정보를 제공하지 않음
- WebQuestions: 하나의 질문과 확장된 관련 답변들로 구성됨.
- CuratedTrec: 실제 사용자들로부터 추출한 쿼리-답변 쌍.
4.4. Main results
세 벤치마크 데이터셋 모두에서 REAML이 성능이 우수. 또한 우수한 성능의 T5 모델과 비교해봤을 때, REALM은 파라미터 수도 적은데 성능도 더 우수함.
*5개의 문서 정도만 검색해서 최고의 성능을 달성
