[논문 리뷰] Data Selection for Language Models via Importance Resampling

Sang Michael Xie, Shibani Santurkar외 2명|arXiv (Cornell University)|2023. 02. 06.

Natural Language Processing Techniques인용 수 16

한 줄 요약

본 논문은 DSIR을 도입합니다. DSIR은 축소된 특성 공간(해시된 n-그램)에서 중요도 가중치를 추정하고 이를 재샘플링하여 목표 분포와 일치시키는 확장 가능한 프레임워크로, 무작위 또는 휴리스틱 베이스라인 대비 다운스트림 성능을 향상시킵니다.

ABSTRACT

Selecting a suitable pretraining dataset is crucial for both general-domain (e.g., GPT-3) and domain-specific (e.g., Codex) language models (LMs). We formalize this problem as selecting a subset of a large raw unlabeled dataset to match a desired target distribution given unlabeled target samples. Due to the scale and dimensionality of the raw text data, existing methods use simple heuristics or require human experts to manually curate data. Instead, we extend the classic importance resampling approach used in low-dimensions for LM data selection. We propose Data Selection with Importance Resampling (DSIR), an efficient and scalable framework that estimates importance weights in a reduced feature space for tractability and selects data with importance resampling according to these weights. We instantiate the DSIR framework with hashed n-gram features for efficiency, enabling the selection of 100M documents from the full Pile dataset in 4.5 hours. To measure whether hashed n-gram features preserve the aspects of the data that are relevant to the target, we define KL reduction, a data metric that measures the proximity between the selected pretraining data and the target on some feature space. Across 8 data selection methods (including expert selection), KL reduction on hashed n-gram features highly correlates with average downstream accuracy (r=0.82). When selecting data for continued pretraining on a specific domain, DSIR performs comparably to expert curation across 8 target distributions. When pretraining general-domain models (target is Wikipedia and books), DSIR improves over random selection and heuristic filtering baselines by 2-2.5% on the GLUE benchmark. Code is available at https://github.com/p-lambda/dsir.

연구 동기 및 목표

• 대규모 라벨이 없는 데이터셋의 부분집합을 선택하여 목표 분포와 일치시키는 데이터 선택의 형식을 formalize한다
• 실용적인 특성 공간에서 중요한 가중치를 추정하는 확장 가능한 방법을 개발한다
• 해시된 n-그램 특성이 데이터 선택에 필요한 목표 관련 정보를 보존하는지 입증한다
• DSIR의 도메인 특화 지속적 사전학습 및 일반 도메인 LM 사전학습에서의 효과를 보여준다
• 특징 공간 재현성을 다운스트림 성능과 연결하는 지표로 KL 축소를 도입한다

제안 방법

• h를 통해 원시 데이터와 타깃 데이터를 특성 공간으로 매핑하고 q_feat와 p_feat를 생성한다
• 각 원시 샘플 z_i = h(x_i)에 대해 중요도 가중치 w_i = p_feat(z_i) / q_feat(z_i)를 추정한다
• 정규화된 가중치에 따라 중복 없이 원시 데이터에서 k개의 샘플을 재샘플링한다
• 해시된 n-그램 특징(m 버킷, 예: m=10000)과 bag-of-ngrams 모델 매개변수 gamma로 DSIR을 구현한다
• 해시 버킷 빈도수를 세어 특징 분포를 추정하고 gamma_hat를 얻는다
• The Pile에서 1 CPU 노드로 4.5시간 만에 1억 개의 문서를 선택하는 확장성을 시연하고, 타깃 유사성의 proxy로 KL 축소를 측정한다

실험 결과

연구 질문

• RQ1 축소된 특성 공간에서의 중요도 재샘플링이 목표 분포와 일치하는 사전 학습 데이터를 효과적으로 선택할 수 있는가?
• RQ2 해시된 n-그램 특징이 다운스트림 목표에 관련된 데이터를 선택하는 데 필요한 정보를 보존하는가?
• RQ3 도메인 특화 지속적 사전학습 및 일반 도메인 사전학습에서 DSIR이 무작위 선택 및 휴리스틱 필터링에 비해 다운스트림 성능을 향상시키는가?
• RQ4 데이터 선택 방법 간에 KL 축소가 다운스트림 성능의 신뢰할 수한 저비용 프록시인가?
• RQ5 한 타깃에 대해 데이터를 선택하고 다른 타깃에서 평가할 때 타깃 도메인 정렬이 전이(transfers)에 어떤 영향을 미치는가?

주요 결과

• DSIR은 도메인 특화 지속적 사전학습에서 무작위 선택 및 휴리스틱 베이스라인 대비 평균 약 2%의 향상(F1, 8개 데이터셋 전체 평균)이며 수작업 큐레이션에 비견됩니다.
• 일반 도메인 사전학습에서 DSIR은 베이스라인 대비 GLUE 성능을 평균 2–2.5% 향상시킵니다.
• 해시된 n-그램 특징으로 The Pile에서 1억 개의 문서를 4.5시간 내에 선택할 수 있으며, 대부분의 시간은 중요도 가중치 계산에 소요됩니다.
• 해시된 n-그램에서의 KL 축소가 다운스트림 성능과 강한 상관관계를 보이며 (Pearson r = 0.82)
• 판별적(Discriminative) DSIR 변형은 생성형 DSIR보다 성능이 낮고; unigram+bigram(해시된 버전)이 unigram 전용 버전보다 우수합니다.
• DSIR은 자동 데이터 선택을 가능하게 하여 많은 설정에서 수작업 큐레이션을 대체할 수 있으며, 품질 필터와 계산 예산이 결과를 좌우합니다.
• 도메인 내 사전학습 전이가 DSIR로 데이터를 선택할 때 도메인 간 전이보다 더 잘 일어나며(평균 약 1.7% 개선)