Data Selection for Language Models via Importance Resampling

Sang Michael Xie,Shibani Santurkar,Tengyu Ma,Percy Liang

2023-11-19

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 <a class="link-external link-https" href="https://github.com/p-lambda/dsir" rel="external noopener nofollow">this https URL</a>.

Computation and Language,Machine Learning

What problem does this paper attempt to address?

The problem that this paper attempts to solve is how to select a subset to match the target distribution given a large amount of original unlabeled data sets. Specifically, the paper focuses on the data selection problem of language models, that is, selecting a subset from a large and diverse original data set (such as The Pile data set) so that its distribution in some feature spaces is similar to the target distribution. This problem is crucial for pre - training language models (whether they are general - domain or specific - domain language models), because selecting appropriate pre - training data can significantly improve model performance.

The paper proposes a framework named Data Selection with Importance Resampling (DSIR). This framework selects data by estimating importance weights in the reduced feature space and performing importance resampling according to these weights. This method aims to overcome the statistical infeasibility of estimating importance weights on high - dimensional text data while maintaining computational efficiency and scalability. By using hashed n - gram features, DSIR can efficiently select 100 million documents from The Pile data set in only 4.5 hours. In addition, the paper defines a data metric called KL reduction, which is used to evaluate the degree of distance reduction between the selected data set and the target in some feature spaces. The results show that KL reduction is highly correlated with the average performance of downstream tasks (Pearson correlation coefficient \(r = 0.82\)).