RETRO: Improving language models by retrieving from trillions of tokens

1. 개요

• Trillion단위의 토큰 데이터 베이스로 구성된 retrieval system을 이용한 언어 모델.
• Retrieval-Enhanced Transformer(RETRO)를 사용.
• 25배 적은 파라미터(7B)로 GPT-3와 Jurassic-1과 비슷한 성능을 얻었다.
• Retro는 frozen BERT를 이용한 retriever, cross attention을 이용한 encoder, chunked cross-attention들을 통해 구성되었다.
• 처음부터 학습해도 되고, 기존에 있는 모델들을 retrieval과 결합(Retrofit)해 사용할수 있다.

2. RETRO

bpb(bits-per-bytes, bit per UTF-8 encoded byte)
the pile 논문에서 제안된 metric으로 일반적으로 서로 다른 tokenization 체계에서 perplexity를 사용하기 어려울 때 비교를 위해 bpb를 사용한다.

$BPB = (L_T/L_B)log_2(e^l) = (L_T/L_B)l/ln(2)$
$L_T$: 토큰 단위의 데이터셋의 길이
$L_B$: UTF-8 Encoded byte 단위의 데이터셋 길이
$l$: negative logg likelihood loss

• Retro 사용시 bpb가 더 낮음
• 모델 크기가 증가할때 bpb 낮아짐.
• RETRO가 모델 생성 능력 향상에 도움을 준다.

2.1. 데이터

• MassiveText
  • 5 trilliion token
    

2.2. 토크나이저

• sentencepiece, 128K Token Vocab

2.3. 구조

• Frozen kNN Retriever
  • key-value로 구성된 database.
    • key: frozen bert embedding
    • value: text
  • 들어온 입력을 chunk 단위로 나누고, 각 chunk 별로 database에서 k-neighbour 추출
    $L_2\ distance\ of\ BERT embedding\ d(C,N) = ||BERT(C) - BERT(N)||^2_2$
  • 실제 계산은 SCaNN Library 사용
    
    
• Transformer Encoder
  • 기존 encoder에 추가로 retrieval에서 나온 k neighbour와 입력들에 대해 cross attention
• RETRO Block
  • Causal Attention
  • Chunked Cross Attention
  • FFW

2.4. 의사코드

비교 및 결과

Retrieval

기존 retrieval 간 접근 방법 비교.

Scaling

모델 사이즈 증가에 따른 성능 비교

The Pile

The Pile 데이터에서 세부 항목별 Jurassic, Gopher와의 비교. 많은 부분에서 비슷하거나 더 좋은 성능을 보임.

하지만 몇몇 데이터에서는 부족한 성능

ubuntu_irc

• ubuntu_irc: Ubuntu's Internet Relay Chat technical support channels.
• reorganized the file structure, removed some unnecessary system messages, and compressed them to make it easier to obtain
• ubuntu irc 데이터에서 성능이 비교 모델 대비 낮음.
  

Wikitext103

• 베이스라인 트랜스포머, kNN-LM과의 비교
  

QA

• DPR이나 RAG에 비해 성능은 좋지만, FiD 비해서는 부족한 성능을 보인다.
• 이유는 FiD의 T5가 encoder의 입력을 보다 잘 사용하는 것으로 생각.