SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, Percy Liang
{xie, shibani, tengyuma, pliang}@cs.stanford.edu
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. $^{1}$
Given a fixed compute budget, the choice of pretraining data is critical for the performance of language models (LMs) [18; 10; 24; 64; 27]. Existing works rely on heuristics to select training data. For example, GPT-3 [10] and PaLM [14] filter web data for examples that are closer to formal text from Wikipedia and books as a proxy for high quality, a method which we call heuristic classification. Specifically, they train a binary classifier to discriminate formal text from web data and select web examples that have a predicted probability above a noisy threshold [10; 18; 21]. However, heuristic classification does not guarantee that the selected data is distributed like formal text. As a second example, domain-specific LMs such as Minerva [45] and Codex [11] (math and code LMs, respectively) employ domain-adaptive pretraining (DAPT) [24], where the model is initialized from a base LM and continues to be pretrained on a domain-specific dataset to achieve gains over the base LM on that domain. The domain-specific datasets are typically manually curated, but a framework for automating data selection could save effort and increase the amount of relevant training data.
$^{1}$ Code, selected data, and pretrained models are available at https://github.com/p-lambda/dsir.
Figure 1: Given a large raw dataset such as The Pile [21] and a smaller target dataset (e.g., Wikipedia + books), we aim to select a subset of the raw data that is distributed like the target in some feature space. Our method, DSIR, first estimates importance weights using raw and target data in an n-gram feature space. The importance weights are used to resample a subset of the raw dataset.
In this paper, we consider the problem of data selection: given a large and diverse raw dataset (e.g., The Pile [21]) and a smaller dataset sampled from a desired target distribution, choose a subset of the raw data that is distributed similarly to the target (Figure 1). While a natural approach is to resample the raw data according to importance weights (importance resampling [68]), estimating importance weights on high dimensional data such as text is often statistically intractable [6; 75].
Instead, our Data Selection with Importance Resampling (DSIR) framework efficiently estimates importance weights over a featurization of the raw and target distributions (Section 3). Our framework first maps the raw and target data onto some feature space and resamples a subset of raw data according to importance weights computed in this feature space. DSIR is extensible via the choice of feature space and importance estimator, which specify what aspects of the data we care about.
What is a feature space that both allows for efficient computation and preserves aspects of the data that are relevant for the target? In Section 4, we instantiate DSIR with hashed n-gram features, where n-grams are hashed onto a fixed number of buckets, for efficiency and scalability. The importance estimator is parameterized by bag-of-words generative models on the hashed n-grams, learned by simply counting the hash bucket frequencies. DSIR with hashed n-grams enables the selection of 100M documents from the full Pile dataset in 4.5 hours.
To evaluate how well hashed n-grams preserves the aspects of the data that are relevant for the target, in Section 6 we define KL reduction, a data metric that measures how much a selected dataset reduces the Kullback-Leibler (KL) divergence to the target (in some feature space) over random selection (KL(target||random) - KL(target||selected)). We show in Section 5 that KL reduction highly correlates with average downstream performance (Pearson  $r = 0.82$ ) across 8 data selection methods, including expert selection.
We consider selecting data from The Pile (1.6B examples) for continued pretraining of domain-specific LMs and training general-domain LMs from scratch. First, we select data for continued pretraining [24] of domain-specific LMs in a controlled setting where the target samples are unlabeled training inputs from a downstream dataset (Section 5). We perform continued pretraining on the selected data starting from RoBERTa [48] and evaluate by fine-tuning on the downstream dataset (whose unlabeled inputs were also used as the target for data selection). On 8 datasets from 4 domains (CS papers, biomedical papers, news, reviews), DSIR improves over RoBERTa (no continued pretraining) by  $2\%$  on average and is even comparable to continued pretraining on expert-curated data from Gururangan et al. [24].
For general-domain LMs (Section 7), the data selection target is formal, clean text from Wikipedia and books, following GPT-3 [10]. We train a masked language model (MLM) from scratch on the selected data and evaluate by fine-tuning on GLUE [82]. In controlled experiments, heuristic classification performs comparably to random sampling from The Pile, possibly because The Pile is already filtered
using heuristic classification. DSIR improves over both baselines by  $2 - 2.5\%$  on average on GLUE. We publish the selected dataset and the code for DSIR to improve future LM pretraining.
Given a small number of target text examples  $x_1', x_2', \ldots, x_n'$  from a target distribution of interest  $p$  and a large raw dataset  $x_1, x_2, \ldots, x_N$  drawn from distribution  $q$ , we aim to select  $k$  examples ( $k \ll N$ ) from the raw dataset that are similar to the target.
Selection via heuristic classification. As a starting point, we first define the heuristic classification method used by GPT-3/The Pile/PaLM [10; 21; 14]. In heuristic classification, we train a binary classifier  $f: \mathcal{X} \to [0,1]$  to output the probability that an input is sampled from the target distribution. The model is typically a fasttext linear classifier on n-gram feature vectors (usually unigrams and bigrams) [30]. We initialize the feature vectors from pretrained fasttext word vectors. We use the trained classifier to estimate  $f(x_{i})$ , the predicted probability that  $x_{i}$  is sampled from the target, for all raw examples. Example  $x_{i}$  is selected if  $f(x_{i}) > 1 - \beta_{i}$ , where  $\beta_{i}$  is a sample from a Pareto distribution (typically with shape parameter  $\alpha = 9$  [10; 14]). Since each example is kept or discarded independently, to select a desired number of examples  $k$ , the process must either be repeated or  $\alpha$  must be tuned. Heuristic classification selects examples from modes of the target distribution (high  $f(x_{i})$ ), which could lack diversity. To combat this, noise  $\beta_{i}$  is added. However, it is unclear how much noise to add and there are no guarantees on the selected data distribution.
3 Data Selection with Importance Resampling
In the DSIR framework, we consider using importance resampling [68] to select examples that are distributed like the target. However, estimating importance weights on high dimensional data like text is often statistically intractable without sufficient additional structure [6; 75; 23].
Instead, we employ importance resampling on a feature space  $\mathcal{Z}$  that provides this structure. DSIR uses a feature extractor  $h: \mathcal{X} \to \mathcal{Z}$  to map the input  $x$  to features  $z = h(x)$ . The induced raw and target feature distributions are  $q_{\mathrm{feat}}$  and  $p_{\mathrm{feat}}$ , respectively. The goal is to select examples with features that are approximately distributed according to the target feature distribution  $p_{\mathrm{feat}}$ . Depending on the choice of feature extractor, DSIR focuses on different aspects of the input. For example, an n-gram feature extractor focuses on matching the n-gram frequencies of the selected data and the target.
Our framework consists of 3 steps:
1. Learn  $\hat{p}_{\text {feat }}$  and  $\hat{q}_{\text {feat }}$ : We learn two feature distributions  $\hat{p}_{\text {feat }}$  and  $\hat{q}_{\text {feat }}$  using held-out featurized examples from the target and raw data, respectively.
2. Compute importance weights: We compute the importance weights  $w_{i} = \frac{\hat{p}_{\mathrm{feat}}(z_{i})}{\hat{q}_{\mathrm{feat}}(z_{i})}$  for each featurized example  $z_{i} = h(x_{i})$  from the  $N$  raw examples.
3. Resample: Sample  $k$  examples without replacement from a categorical distribution with probabilities  $\frac{w_i}{\sum_{i=1}^{N} w_i}$ . Sampling without replacement avoids choosing the same example multiple times, is more statistically efficient for importance resampling [23], and can be implemented efficiently with the Gumbel top- $k$  trick [81; 38; 87; 39].
4 DSIR with Hashed N-gram Features
For efficiency and scalability, we instantiate DSIR with hashed n-gram features. Later in Section 6, we test whether hashed n-gram features preserve the information needed to select data relevant to the target.
Hashed n-gram features. We consider hashed n-gram features inspired from fasttext [30; 86]. Specifically, for each example  $x$  we form a list of unigrams and bigrams, hash each n-grams into one of  $m$  buckets ( $m = 10000$  in this paper), and return the counts of the hashed buckets in a  $m$ -dimensional feature vector  $z \in \mathbb{N}^m$ . For example, if the text input is "Alice is eating", we form the list [Alice, is, eating, Alice is, is eating], hash each element of this list to get a list of indices [1, 3, 3, 2, 0] and return the vector of counts for each index [1, 1, 1, 2, ...]. While the hashing
introduces some noise due to collisions, we find that this is a simple and effective way to incorporate both unigram and bigram information.
Bag of hashed n-grams model. We parameterize the raw and target feature distributions  $p_{\mathrm{feat}}$  and  $q_{\mathrm{feat}}$  as bag-of-ngrams models. The bag-of-ngrams model has parameters  $\gamma \in \Delta^m$ , which is a vector of probabilities on the hash buckets that sums to 1. Under this model, the probability of a feature vector  $z \in \mathbb{N}^m$  is
$$
\mathbb {P} (z; \gamma) = \prod_ {j = 1} ^ {m} \gamma [ j ] ^ {z [ j ]} \tag {1}
$$
where the bracket notation selects the corresponding index in the vector. Given some featurized examples  $\tilde{z}_1,\ldots ,\tilde{z}_s$  sampled from a feature distribution, we estimate the parameters by counting:  $\hat{\gamma} = \frac{1}{\sum_{i = 1}^{s}\mathbf{1}^{\top}\tilde{z}_{i}}\sum_{j = 1}^{s}\tilde{z}_{j}$
Speed benchmark on The Pile. To test the scalability of the framework, we benchmark DSIR with hashed n-gram features on selecting data from the full Pile dataset [21]. For this test, we do not preprocess the data other than decompressing the text files for faster I/O. We use hashed n-gram features with 10k buckets, fit the raw feature distribution with 1B hashed indices from the Pile, and fit the target feature distribution with the full target dataset (ChemProt [41]). DSIR selects 100M documents from the full Pile dataset in 4.5 hours on 1 CPU node with 96 cores. Almost all of the time (4.36 hours) is spent computing the importance weights on the raw dataset, while fitting the feature distributions (1 minutes) and resampling (6 minutes) were much faster. Increasing the number of CPU cores can further decrease the runtime.
5 Selecting Data for Domain-Specific Continued Pretraining
In this section, we use DSIR to select domain-specific data for continued pretraining. We compare DSIR to 7 other data selection methods in this continued pretraining setting.
Setup. We select data for 8 target distributions in the setting of Gururangan et al. [24], where we perform continued pretraining of domain-specific LMs. Here, the target is a specific downstream unlabeled data distribution and we select examples from The Pile (the raw data). For each downstream dataset, we select data for continued pretraining starting from RoBERTa [48] (see Appendix H). Following Gururangan et al. [24], we consider 8 downstream datasets across 4 domains: Computer Science papers (ACL-ARC [31], Sci-ERC [51]), Biomedicine (ChemProt [41], RCT [16]) News (AGNews [92], HyperPartisan [34]), and Reviews (Helpfulness [53], IMDB [52]).
Baselines. Beyond random selection (without replacement) and heuristic classification, we also compare against manual curation [24] and a top- $k$  variant of heuristic classification. In manual curation, we simply fine-tune from domain-adaptive pretraining (DAPT) checkpoints [24], which are the result of continued pretraining on manually-curated data. In top- $k$  heuristic classification, we select the top- $k$ -scoring examples according to the binary classifier used in heuristic classification. All methods select data from The Pile except for manual curation, which uses domain-specific data sources [24].
We perform a controlled comparison by equalizing the amount of LM training compute for all methods, measured by the number of tokens processed during training, following the compute budget in Gururangan et al. [24]. For random selection, heuristic classification, and DSIR using n-gram features (defined in Section 4), we control the number of selected examples (25M examples with fixed token length 256) and the training protocol. We standardize the fine-tuning for all models and average all results over 5 random seeds (see Appendix H for details). All the models initialize from RoBERTa-base. Before data selection via DSIR or heuristic classification, we remove extremely short (<40 words) or repetitive documents that tend to be uninformative (Appendix J).
Automatic data selection with DSIR can replace manual curation. Table 1 shows the comparison between the data selection methods. To summarize:
- On average, DSIR improves over random selection by  $1.2\%$  and manually curated data (DAPT) by  $0.3\%$ , showing the potential to replace manual curation.
Table 1: F1 scores for continued pretraining from the RoBERTa checkpoint [48] on 8 downstream datasets from 4 domains (CS, Biomed, News, and Reviews). Random selection, heuristic classification, and DSIR train on 25M selected examples from The Pile. Heuristic classification and DSIR create a different pretraining dataset for every downstream dataset. All models (including DAPT [24]) use the same amount of training compute and results are averaged over 5 seeds, with standard deviations in subscripts. All datasets use macro-F1 except ChemProt and RCT, which use micro-F1.
ACL-ARC	Sci-ERC	ChemProt	RCT	HyperPartisan	AGNews	Helpfulness	IMDB	Avg
RoBERTa (no continued pretrain)	66.801.08	80.142.25	82.310.54	86.680.14	88.852.59	93.350.2	65.082.29	94.380.13	82.20
Random selection	67.512.60	80.531.65	83.140.52	86.850.13	86.425.33	93.520.15	68.151.37	94.490.25	82.58
Manual curation/DAPT [24]	71.844.78	80.421.57	84.170.50	87.110.10	87.233.65	93.610.12	68.211.07	95.080.11	83.46
Heuristic classification	69.942.96	80.520.95	83.351.07	86.780.17	85.716.01	93.540.19	68.500.79	94.660.22	82.88
Top-k Heuristic classification	71.730.21	80.220.58	84.110.73	87.080.21	88.298.28	93.670.14	69.180.73	94.900.14	83.65
DSIR	72.862.71	80.441.13	85.510.46	87.140.13	87.014.53	93.620.19	68.950.81	94.560.34	83.76
Table 2: F1 scores for continued pretraining from RoBERTa, testing different components of heuristic classification and DSIR methods. For heuristic classification, replacing the Pareto noisy threshold with calibration and importance resampling (DSIR (fasttext discriminative)) improves F1. Generative importance weight estimators outperform discriminative importance weight estimators for DSIR. All results average over 5 seeds, with standard deviations in subscripts.
ACL-ARC	Sci-ERC	ChemProt	RCT	HyperPartisan	AGNews	Helpfulness	IMDB	Avg
Heuristic classification	69.942.96	80.520.95	83.351.07	86.780.17	85.716.01	93.540.19	68.500.79	94.660.22	82.88
DSIR (n-gram generative)	72.862.71	80.441.13	85.510.46	87.140.13	87.014.53	93.620.19	68.950.81	94.560.34	83.76
DSIR (fasttext discriminative)	68.467.15	79.001.50	84.570.65	87.090.08	89.184.06	93.540.14	68.411.51	94.950.29	83.15
DSIR (n-gram discriminative)	70.352.90	80.210.85	85.031.18	87.040.19	85.498.24	93.740.07	68.791.22	94.840.24	83.19
DSIR (unigram generative)	69.530.16	79.691.91	85.240.88	87.050.10	90.115.39	93.420.16	68.550.78	94.390.33	83.50
- DSIR improves over heuristic classification by  $0.9\%$  and is comparable to top- $k$  heuristic classification. We note that top- $k$  heuristic classification is not typically used in this setting, but we find that it may be particularly suited for domain-specific data selection, where diversity may be less important than the general-domain setting.
- Random selection improves by  $0.4\%$  on average over no continued pretraining at all, showing that additional data generally improves the downstream F1 score. All the targeted data selection methods improve over random selection.
Discriminative importance weight estimators underperform generative estimators. We experiment with replacing components of DSIR in Table 2. First, we consider using the binary classifier from heuristic classification (which takes pretrained fasttext word vectors as input) as the importance weight estimator in DSIR. For input  $x_{i}$ , the classifier predicts the probability of target  $f(x_{i})$ . We use this to estimate importance weights  $\frac{f(x_i)}{1 - f(x_i)}$ , then resample according to these weights. This approach (DSIR (fasttext discriminative)) improves F1 by  $0.3\%$  over heuristic classification. However, this approach still underperforms DSIR by  $0.6\%$  on average, even with regularization and calibration.
We consider another discriminative version of DSIR that uses hashed n-gram features as input to a logistic regression binary classifier for importance weight estimation. This differs from heuristic classification, which initializes with pretrained fasttext feature vectors and fine-tunes the features along with the classifier. This approach (DSIR (n-gram discriminative)) underperforms DSIR by  $0.7\%$ , even with regularization and calibration. These results suggest that a generative approach is better suited (or easier to tune) for importance resampling. However, the discriminative approaches still outperform random selection by  $0.6\%$ .
Selecting with n-grams improves downstream performance over unigrams. DSIR uses both unigram and bigram information to compute hashed n-gram features. We ablate the role of bigram information in hashed n-grams by using hashed unigram features (with 10k buckets) for DSIR. In Table 2, we find that DSIR with unigram features underperforms DSIR with n-grams by  $0.26\%$ , though still achieving comparable F1 score to manual curation. Overall, selecting data with unigrams is effective, but including bigrams further improves the relevance of the selected data.
Cross-domain analysis and the effect of the choice of pretraining data. DSIR assumes knowledge of the target distribution, but what happens if the target dataset is not representative of the target
Figure 2: F1 scores of DSIR for all pairs of pretraining data target distributions (rows) and downstream tasks (columns). The cells are colored by its per-column ranking, with better rankings (higher F1 scores) having darker colors. While using the pretraining data selected specifically for the downstream task is typically strong, choosing the worst pretraining dataset for the downstream task reduces F1 by  $6\%$  on average. All results are averaged over 5 seeds.
distribution? To test the effect of varying the target distribution on downstream performance, we consider every pair of pretraining dataset, which is selected by DSIR for target downstream task X, and downstream task Y. Figure 2 provides the full matrix of results. We find a  $6\%$  average drop in F1 when we choose the worst pairing for each downstream task instead of matching the pretraining and downstream data. In the worst case, the F1-score on HyperPartisan drops by  $30\%$ . Thus, the choice of target distribution can have a large effect on downstream performance.
Pretraining data transfers better for targets within the same domain. In practice, we may have access to some target datasets in the relevant domain and hope to select pretraining data that can improve performance on other tasks in that domain. The 8 target/downstream tasks we use come from 4 domains, with 2 tasks from each domain. We define within-domain F1 as the average F1 of the pairs of pretraining and fine-tuning data from the same domain, but excluding pairs where the pretraining data is selected for the fine-tuning task. We compute this by averaging the off-diagonal elements in the  $2 \times 2$  diagonal blocks of the matrix in Figure 2. We find that the within-domain F1 (82.9%) is 1.7% higher on average than the cross-domain F1 (81.2%), where the pretraining data is selected for a target from a different domain.
6 KL Reduction on Hashed N-grams Predicts Downstream Performance
When designing a feature extractor for DSIR, how do we measure whether the features preserve the information for selecting relevant pretraining data? To answer this question, we propose a data metric, KL reduction, which measures how much data selection reduces distance to the target over random selection in a feature space. We find that KL reduction on hashed n-gram features strongly correlates with downstream performance across various data selection methods, including those that do not involve n-grams, such as manual curation.
Figure 3: Plot of average KL reduction on the n-gram feature space, defined as how much the selected dataset reduces KL divergence to the target distribution over just random sampling from The Pile, against average downstream F1 score over the 8 continued pretraining datasets in Table 1. There is a strong correlation between KL reduction and downstream performance (Pearson  $r = 0.82$ ).
KL reduction metric. We define KL reduction as the average reduction in empirical KL divergence from doing data selection over random selection over a set of target feature distributions  $\mathcal{T}$ :
$$
\operatorname {K L} - \operatorname {r e d u c t i o n} \left(p _ {\text {f e a t}} ^ {\prime}; \hat {q} _ {\text {f e a t}}, \mathcal {T}\right) = \frac {1}{| \mathcal {T} |} \sum_ {\hat {p} _ {\text {f e a t}} \in \mathcal {T}} \operatorname {K L} \left(\hat {p} _ {\text {f e a t}} \| \hat {q} _ {\text {f e a t}}\right) - \operatorname {K L} \left(\hat {p} _ {\text {f e a t}} \| p _ {\text {f e a t}} ^ {\prime}\right) \tag {2}
$$