[Daily] CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training

지금 하고있는 연구랑 비슷해서 좀 자세히 읽어봤당 역시 OOD 재밌다

TLDR;

• Web data는 보통 web에서 수집되기 때문에 explicit domain label이 없는데, domain-specific training을 위해 optimal pre-training data mixture를 identify하는건 어려운 문제임.
• Cluster-based로 최적의 data mixture weight를 도출하는 framework -> Efficient domain-specific pre-training

Motivation

• Domain-specific task의 성능을 올리는데는 final pre-training phase가 중요하다고 함.
• General/ domain-specific task에 맞는 pre-training data mixture를 도출하는건 어려운 문제임.
  • 예를 들어 coding task를 푸는데는 mathematics, reasoning, and security 등의 complementary knowledge가 쓰일 수 있다.

Method

• Efficient domain-specific training을 가능하게 하는 data mixture optimization 방법 제안

1. Clustering

• 각 raw dataset에 대해 embedding model로 text embedding을 추출한 뒤, kmeans로 clustering
• Pruning: low-quality cluster를 지움
• Merge: 유클리디안 거리 기반으로 similar한 cluster를 merge

2. Iterative Bootstrapping: Mixture Weight Search

• Sampling mixture weight (domain-specific cluster를 어떻게 섞을까 인 것 같음)를 optimize 하기 위한 과정

• Bi-level optimization problem으로 푼다
  • 궁극적으론 task performance를 최대화 하는 mixture weight $\alpha$를 구하는 것
  • 근데 이렇게 모든 alpha configuration을 시도해서 task performance를 보기엔 학습을 일일이 다 해야해서 너무 time-consuming 함
• 따라서 task performance를 approximate 할 수 있는 proxy model을 학습시킨다.
  • 이 proxy model은 아무 regression model이나 가능하다.
  • 결국 이 proxy model을 이용해 특정 configuration의 performance를 approximate 한 후 task performance를 maximize하는 적절한 alpha (weight 조합)을 찾는게 목표이다.

• Iterative하게 bootstrap 한다
  • 랜덤으로 initialize한 주어진 configuration 중에서 성능이 가장 높은 top N을 고르고, 이를 이용해 predictor를 학습해 alpha를 도출한다. 이 과정을 반복함. (Iteration)

Experiments

• Domain specific한 밴치마크와 general한 밴치마크에서 모두 좋은 성능을 보인다.

• Weight analysis -> iteration을 거듭할수록 weight이 modify 된다.