[논문 리뷰]IN-CONTEXT PRETRAINING: LANGUAGE MODELING BEYOND DOCUMENT BOUNDARIES

Large Language Model 을 사전학습 하기 위해서는 long input context 가 필요하다.
최근 모델들은 대부분 8K 를 사용하는데 8K 만큼의 token 을 가진 Documents 는 적다.
그래서 보편적인 방법으로 Document 를 이어붙여 8K 만큼 만들어준 뒤 모델을 학습시키는 방법을 사용한다.
본 논문에서는 Document 를 어떻게 이어붙여야 성능 향상의 효과가 있는지를 실험한 내용이다.

목차

1. 논문 핵심 내용 요약

2. 장점 & 단점 정리

3. 코드 유무 확인

---

 

1. 논문 핵심 내용 요약

LLM 학습 시 input contexts 구성할 때 문서의 type 을 고려하여 이어 붙여 모델을 학습하면 성능 향상 효과가 있다고 한다.
본 논문에서는 수십억 개의 문서에서 반복 없이 모든 문서에 대해 문맥 유사성을 극대화하여 정렬하는 알고리즘을 통해 모델을 학습했을 때의 효과를 보여준다.

 

문제 제기

• 기존 사전학습 시에는 짧은 문서 세트를 무작위로 연결하여 input context 로 구성하여 학습
  • 이전 문서는 다음 문서를 예측하는데 적합하지 않은 경우가 많음.
  • 다양한 sequence 길이로 사전학습학습할 때 학습 데이터를 여러 bucket(ex. <2K, 16K, 64K)으로 나누고 배치당 context length 를 조정하며 학습하기도 함
• 수십억 개의 문서에서 일련의 관련 문서를 이어붙이는 작업은 어려움

 

해결 방안 및 결과

• 일련의 관련 문서를 이어붙여 input contexts 로 구성하여 사전학습 한다.
  • 인련의 관련 문서를 구하기 위해 efficient nearest neighbor search 활용.
  • graph traversal algorithm 으로 일관된 input contexts 를 구성.
• 결과: LLM 의 성능이 향상
  • 복잡한 문맥 추론이 필요한 작업에서 눈에 띄게 향상된 성능을 보임

 

방법론

2개의 단계로 구성되어 있다.

IN_CONTEXT PRETRAINING 논문 figure

1. 관련된 문서들끼리 묶는다.
   • dense retrieval model 사용하여 문서를 embedding
     • 모델 정보는 Unsupervised dense information retrieval with contrastive learning 논문 참고
   • 문서들마다 top-k 유사도를 구함 (2개의 문서 사이의 유사도) → cos 유사도
   • approximate nearest neighbour search with the faiss library 사용하여 관련 문서 묶기
2. input contexts 를 구조화한다.
   • 관련 문서를 활용하여 input context 구성을 위해 nearest neighbors 기준으로 각 문서를 붙인다.
   • 동일한 문서가 input context 로 활용되는 문제를 피하기 위해 document graph traversal 사용
     • document 를 node 로, document 간 유사도를 node 간의 weight 로 하는 graph 생성
     • weight 가 가장 큰 document 기준으로 input context 로 merge 하고 merge 된 document 는 제거하는 방식으로 진행 → 상세한 내용은 논문 본문의 2.2 참고

 

성능 평가 결과

• Model: LLaMA archi 기준, 0.3, 0.7, 1.5, 7B size 만큼 scratch 부터 학습
  • Context Length: 8192
• Datasets: 2개의 버전을 비교
  • 랜덤으로 docs 를 input contexts 로 구성 vs 관련 docs 끼리 input contexts 로 구성
• 성능 평가 (eval Datasets & Metrics 은 논문 3.3.1~ 3.3.6 참고)
  • including in-context learning (+8%)
  • reading comprehension (+15%)
  • faithfulness to previous contexts (+16%)
  • long-context reasoning (+5%)
  • retrieval augmentation (+9%)

---

 

2. 장점 & 단점 정리

장점

• LLM Pretrain 시 input context 구성에 대한 비교 실험이 유의미하다고 생각한다.
  
  • 통제해야 하는 변수를 명확히 설정했기 때문.
• 유사도가 높은 docs 를 input contexts 로 구성했을 때 유의미한 성능 향상이 있었다.

 

단점

• 관련 docs 끼리 군집화 시키는 성능이 우려된다.
• 계산량이 많아 시간이 오래 걸린다.
• 동일 docs 가 오직 한 번만 사용되도록 구성하는게 어렵다.

---

 

3. 코드 유무 확인

논문에서 참조한 github 코드이다.

https://github.com/swj0419/in-context-pretraining

GitHub - swj0419/in-context-pretraining

---

 

마무리,,

많은 Reference 들이 SFT, Instruct Tuning 학습 방법과 데이터셋 구성에 focus 되어 있어 Continual Pretrain method 에 대한 Reference 는 많지 않다. 그만큼 어려운 task 라는 의미이기도 한 것 같다.

본 논문은 Continual Pretrain method 중 가장 중요한 datasets 구성에 대해 다루고 있다는 점 자체가 유의미하다고 본다.

그렇기에 저자의 많은 고민이 보이는 실험 설계와 실험 결과 분석이 눈에 띄었다.

 

필자는 실험 시 가장 많은 고민을 하는 사항이 가설 설정과 가설 검증을 위한 변수 셋팅이다.

어떤 변수를 fix 해야 되고 어떤 변수를 change 해야 하는지가 명확해야 유의미한 가설 검증을 할 수 있다.

본 논문은 이러한 실험 설계가 잘 되어 있어 저자의 의도가 한 눈에 보였다.

 

만약 LLM 을 CPT 학습하고자 한다면 해당 논문을 참고하여 Input contexts 구성을 해보기 바란다.