Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a single Transformer

Open-domain QA 시스템은 일반적으로 두 단계로 구성됩니다. 대규모 Corpus에서 관련 문서를 효율적으로 검색한 뒤, 선택된 문서를 자세히 읽어 답변을 생성합니다. Retrieval와 Reader는 보통 각각 모델링 되어 번거로운 구현과 End-to-end Training이 불편한 문제점이 있었다. 

이 연구에서는 위 문제를 해결하기 위해, 이러한 설계를 재검토하고 별도의 아키텍처와 훈련을 피하여 Retrieval as Attention (ReAtt)라는 single Transformer를 도입하여, QA 작업의 최종 supervision 기반한 end-to-end training을 수행합니다.

End-to-end의 single 모델이 경쟁력 있는 검색 및 QA 성능을 달성할 수 있음을 처음으로 보여주며, 별도로 훈련된 Retrieval와 Reader의 성능과 비슷하거나 다소 높일 수 있습니다. 더욱, end-to-end adaptation은 감독 및 비감독 설정 모두에서 아웃 오브 도메인 데이터셋에서 성능을 크게 향상시켜, 이 모델이 지식 집약적 작업에 대한 간단하고 적응 가능한 솔루션임을 입증한다고 합니다.

 

자세한 설명은 다음과 같습니다.

Introduction

question answering (QA), fact checking, and dialogue generation과 같은 Knowledge-intensive tasks는 모델이 방대한 지식 말뭉치(예: 위키백과)에서 관련 정보를 수집하고, 수집된 증거를 바탕으로 답변을 생성해야 한다. 이전에는 바이 인코더 아키텍처를 사용하여 효율성을 위해 query와 document를 독립적으로 인코딩한 후, 검색된 문서를 크로스 인코더 아키텍처를 통해 보다 신중하며 확장가능한 방식으로 읽음으로써 문제를 해결했다. 이는 Retrieval와 Reader를 따로 훈련하는 방법이며, 따로 훈련했을 때 훈련이 번거롭다. 또, 두 모델을 하드 코딩으로 결합할 수 있지만, 이 방법은 효과적인 end-to-end 학습과 새로운 Domain에 대한 adaptation을 방해한다.

 

이후 Retrieval와 Reader를 연결하려는 시도가 있었다. (Lee et al., 2019; Guu et al., 2020; Lewis et al., 2020; Sachan et al., 2021; Lee et al., 2021a; Izacard et al., 2022) 하지만 이 연구에서는 완전한 end-to-end 학습 방법이 아니었으며, 이들은 기존의 supervised training된 Dense Retrieval에서 초기화를 요구하거나(Lee et al., 2020), 비싼 unsupervised Retrieval의 pre-training을 워밍업으로 필요했다.(Lee et al., 2019; Guu et al., 2020; Sachan et al., 2021; Lee et al., 2021a; Izacard et al., 2022).

 

본 논문에서는 Open-Domain QA에 초점을 맞추며, 단일 Transformer 내에서 Retrieval 와 Reader 역할을 모두 수행하면서 완전한 end-to-end training 방식으로 경쟁력 있는 성능을 달성한다.

이전 연구의 annotation과 warm-up의 필요성을 없애고, 검색 보강 훈련을 단순화하여 새로운 Domain에 대한 adaptation을 용이하게 만든다. 본 연구에서는 self-attention과 Query-Document 연결하는 것 사이의 관계를 유추하여 single Transformer 내에서 Retrieval 와 Reader 를 유기적으로 융합할 수 있다고 생각한다.

 

구체적으로, 본 연구에서는 T5 Encoder-Decoder 모델을 사용하며, 이를 Retrieval 와 Reader 모두로 사용한다. 첫 번째 B 개의 인코더 레이어를 bi-encoder 로 사용하여 쿼리와 문서를 독립적으로 인코딩하고, 레이어 B + 1에서의 attention 점수(Retrieval attention 으로 표기됨)를 사용하여 관련성 점수를 계산한다.

본 연구에서는 Retrieval를 위한 self-attention을 직접 사용하는 것이 강력한 Retrievals 보다 성능이 떨어진다는 것을 발견했으며, 이는 self-attetion이 Local context 에 대해 사전 훈련되어 전체 말뭉치의 큰 표현 공간에서 관련 정보를 식별하는 데 충분하지 않기 때문이라고 추측한다.

이를 해결하기 위해, 쿼리와 다수의 문서 간의 Retrieval attention을 계산하고, 조정하는 방법을 제안한다. 각 쿼리에 대해 positive 문서와 hard negative(어려운 부정적인) 문서를 포함할 가능성이 있는 close document와, 같은 배치 내의 다른 query의 document 를 random negatives 로 간주하여 Retrieval attention 을 계산한다. Rertrieval attention 은 decoder와 encoder 간의 cross attention(target attention 이라고 표기됨)와의 차이를 최소화하여 조정되며, 이는 답변 생성에 있어 각 문서의 유용성을 나타낸다(Izacard and Grave, 2021a). 결과적으로 Retrieval as Attention 모델(ReAtt)은 QA annotations 를 기반으로 훈련된 single T5 모델로, cross-document 조정을 통해 유용한 문서를 동시에 촉진하는 방법을 학습한다.

 

본 연구는 ReAtt를 Natural Questions 데이터셋(NQ)(Kwiatkowski et al., 2019)을 사용하여 완전한 End-to-end 방식으로 훈련하며, 완전한 annotated 검색기인 ColBERT-NQ(Khattab et al., 2020)와 강력한 QA 모델 FiD(Izacard and Grave, 2021b,a)의 성능과 비슷한거나 좋은 결과를 얻는다. ReAtt의 일반화 및 end-to-end adaptation 능력을 추가로 테스트하기 위해, 우리는 BEIR 벤치마크(Thakur et al., 2021)의 7개 데이터셋에서 제로샷, supervised, unsupervised adaptatin 실험을 수행한다. 모든 설정에서 end-to-end adaptation 은 검색 성능을 대폭 향상시키며, 강력한 검색 적응 및 사전 훈련 방법과 비슷하거나 높은 성능을 달성한다.

Retrieval as Attention(ReAtt)

Formal Definition

Query와 Document 사이의 상호작용이 없는 bi-encoder 아키텍처를 사용한다. $E_{d} = biencoder(d)$ 를 사용하여 Corpus $ d \in D$ 인 document 의 한개 이상의 multiple representations 를 말하며, 마찬가지로 $E_{q} = biencoder(q)$ 를 사용하여 query representations 을 나타낸다. query와 가장 관련있는 top-k의 documents 는 $D^{ret}_{q} = arg topk_{d \in D^{r}}(E_{q}, E_{d})$ 로 검색 되며, 여기서 함수 $r$은 query와 document representations 을 기반으로 관련성을 계산한다. 이 관련성은 queries 와 documents 가 single vector로 인코딩 되면 단순하게 dot product로 나타낼 수 있고 $D^{ret}_{q}$ 는 반환된 문서를 나타낸다.

 

본 논문에서는 encoder-decoder 기반의 생성 QA를 다루는데, 이는 encoder $E_{q,d} = crossencoder(q,d)$ 로 검색된 문서와 queries 를 함꼐 표현하며, decoder $P^{gen}(a|q,d) = P^{gen}(a|E_{q,d})$ 를 통해 답변 $a$를 autoregressive 하게 생성한다.

검색된 multiple documents를 다루기 위해, 각 query-document pair를 독립적으로 encode 하는 FiD를 사용했으며 decoder의 cross-attention $P^{gen}(a|q, D^{ret}_{q}) = P^{gen}(a|E_{q,d_{1}}, E_{q,d_{2}}, \cdots, E_{q,d_{D^{|ret|}_{q}}})$ 을 통해 이 representations 을 융합한다. Negative log likelihood(NLL) $L_{QA}=- \log P^{gen}(a|q, D^{ret}_{q})$를 optimization 으로 사용한다. 

Leveraging Attention for Retrieval

다음으로 query와 documents 간의 self-attention 을 Retrieval scores로 직접 사용하는 방법을 소개한다.

Putting the Retrieval into Transformers