SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $p_{\mathrm{feat}}'$  is the empirical feature distribution of the selected data,  $\hat{p}_{\mathrm{feat}}$  is a empirical target feature distribution,  $\hat{q}_{\mathrm{feat}}$  is the empirical raw feature distribution. KL reduction depends on the raw distribution  $\hat{q}_{\mathrm{feat}}$  and the set of target distributions  $\mathcal{T}$  as hyperparameters. In our continued pretraining setting,  $\hat{q}_{\mathrm{feat}}$  is the feature distribution of the Pile and  $\mathcal{T}$  consists of the feature distributions from the 8 downstream tasks from Section 5.
KL reduction on hashed n-grams predicts downstream performance. We show that when computed on the hashed n-gram feature space, KL reduction of a selected dataset highly correlates with the downstream performance of a model trained on that data. Figure 3 plots KL reduction against average downstream performance over 8 target distributions for 8 data selection methods from The Pile [21], where the distribution parameters are estimated using 100k samples from each dataset. The average downstream F1 score is highly correlated with the KL reduction (Pearson  $r = 0.82$ ). This agrees with the results of Razeghi et al. [65] for in-context learning [10] and extends the preliminary evidence from Gururangan et al. [24] on one selection method that better unigram overlap improves downstream performance. DSIR with hashed n-gram features achieves the highest KL reduction and the best average downstream F1.
While some of the original pretraining datasets for DAPT [24] were not publicly available, we downloaded the public versions as an approximation. Our results suggest that hashed n-gram features preserve most of the information needed for selecting data relevant to the target. Since KL reduction highly correlates with downstream performance and can be cheaply computed without training an LM, KL reduction can be used as a sanity check for future data selection methods.
7 Selecting Data for Training General-Domain LMs
In this section, we consider selecting formal text (as a proxy for high-quality text) for training general-domain LMs from scratch. We use Wikipedia and books as the target distribution.
Baselines and setup. We compare the following methods for selecting data from the Pile: 1) Random selection, 2) Heuristic classification (GPT-3/Pile/PaLM method), and 3) DSIR. As ablations, we consider top- $k$  variants of heuristic classification and DSIR (take the top- $k$  examples according to importance weights instead of resampling). We use each method to select 51.2M examples, which corresponds to 4 epochs with our compute budget. For heuristic classification and DSIR, we select  $96\%$  of the examples from domains excluding Wikipedia and books. This is done to reduce the bias towards selecting data from Wikipedia and books (the target distribution). We choose the other  $4\%$
Table 3: Accuracies on the GLUE [82] dev set for a BERT-style masked language model [17] trained on data selected from The Pile [21]. Following RoBERTa [48], for RTE, STS, and MRPC we fine-tune starting from the MNLI model instead of from scratch. DSIR outperforms heuristic classification (used by GPT-3 and PaLM) and random selection by over  $2\%$  on average. All results are averaged over 5 seeds and standard deviations are in subscripts.
MNLI	QNLI	QQP	RTE	SST-2	MRPC	CoLA	STS-B	Avg
Random selection	82.630.41	86.900.28	89.570.30	67.371.69	90.050.41	87.401.08	49.413.67	88.630.11	80.25
Heuristic classification	82.690.17	85.950.79	89.770.32	68.591.75	88.940.98	86.030.93	48.173.19	88.620.22	79.85
Top-k Heuristic classification	83.340.22	88.620.24	89.890.19	70.040.99	91.150.76	86.371.00	53.023.56	89.300.11	81.47
DSIR	83.070.29	89.110.14	89.800.37	75.092.76	90.480.57	87.700.68	54.001.34	89.170.13	82.30
Top-k DSIR	83.390.06	88.630.38	89.940.17	72.491.29	91.010.79	86.181.12	49.901.10	89.520.21	81.38
uniformly from Wikipedia and books, and did not tune these proportions (Appendix F). We apply a quality filter for extremely short or repetitive examples before heuristic classification and DSIR selection (Appendix J). For each dataset, we perform MLM pretraining for 50k steps with a large batch size (4096) and short token length (128), following Izsak et al. [28]. All the models use the BERT-base architecture [17]. We evaluate the models on the GLUE dev set, averaged over 5 fine-tuning runs [82]. Fine-tuning hyperparameters such as the number of epochs and batch size are fixed for each dataset, following reasonable defaults from the RoBERTa codebase [48].
DSIR qualitatively selects more formal text. Figure 4 shows the beginning characters of 20 random examples selected by random selection, heuristic classification, and DSIR. The random sample contains many code examples that are not similar to text from Wikipedia and books. Heuristic classification seems slightly too diverse, which suggests that the variance of the Pareto distribution added to the classifier scores may be too high. Note that we use the setting of the Pareto shape hyperparameter used in GPT-3 [10]. Qualitatively, DSIR selects the most formal text. By doing importance resampling to match the target distribution, DSIR trades off the relevance and diversity of the selected data automatically.
DSIR improves GLUE performance. Table 3 shows results on the GLUE dev set. DSIR achieves  $82.3\%$  average GLUE accuracy, improving over random selection by  $2\%$  and heuristic classification by  $2.5\%$ . Heuristic classification leads to  $0.5\%$  lower accuracy than random selection from The Pile. We hypothesize this is because The Pile has already been filtered once with heuristic classification.
Resampling outperforms top- $k$  selection. Top- $k$  heuristic classification and top- $k$  DSIR have similar performance across datasets, with some tradeoffs compared to DSIR without top- $k$ . DSIR without top- $k$  is competitive with these variants on all datasets and achieves a  $0.8 - 0.9\%$  higher average over top- $k$  variants. All the top accuracies across datasets are achieved by DSIR or top- $k$  DSIR.
(a) Random selection
(b) Heuristic classification
(c) DSIR
Figure 4: Beginning characters of 20 random examples (each line is a different example) selected by random selection, heuristic classification, and DSIR, where the target is formal text from Wikipedia + books. Qualitatively, DSIR selects more formal text than random selection and heuristic classification.
8 Related Work
Effect of pretraining data on LMs. The pretraining data has a large effect on LM performance. Lee et al. [44]; Hernandez et al. [26] show that deduplicating data improves LMs, and Baevski et al. [3]; Yang et al. [88] compare using a large web corpus versus Wikipedia. Raffel et al. [64] shows that heuristically filtered data (filtering out short and duplicated examples) improves T5 and Du et al. [18] shows that heuristic classification improves downstream few-shot performance for GLaM. We provide extensive controlled experiments comparing the effect of data selection methods on downstream performance.
Retrieval. Yao et al. [89] use keyword-based retrieval (BM25) to select data for semi-supervised learning. In preliminary tests, we found that out of 6.1M documents retrieved by BM25, there were only 1.8M unique documents (70% were exact duplicates). These duplicate examples can hurt performance [44; 26]. Selecting a desired number of unique documents involves oversampling and de-duplication. Instead, we consider top- $k$  heuristic classification, which has similarities to cosine similarity-based retrieval (since heuristic classification uses an inner product score between pretrained word embeddings and a learned class vector) and avoids retrieving repeated examples.
Data selection in classical NLP. Moore-Lewis selection [56; 2; 20] takes the top- $k$  examples in cross-entropy difference between n-gram LMs trained on target and raw data to score examples, which could over-sample examples from the mode of the target distribution. In Section 7, we found that top- $k$  DSIR, which is a form of Moore-Lewis selection with hashed n-gram LMs, underperforms DSIR by  $0.9\%$  on GLUE. DSIR naturally balances diversity and relevance for use in both domain-specific and general-domain cases, since it uses importance resampling to match the target distribution. Feature-space/n-gram discrepancy measures [29; 69; 47] have also been used in selecting data in the domain adaptation setting. Overall, these methods do not consider importance resampling and do not address the gap between pretraining and downstream tasks: pretraining has a different objective to fine-tuning, pretraining uses unlabeled data that is not task-formatted, and the influence of pretraining data is separated from the final model by the fine-tuning step. Beyond the preliminary evidence that unigram similarity metrics are related to downstream performance in Gururangan et al. [24], we show comprehensively and quantitatively on 8 selection methods that despite the pretrain-downstream gap, n-gram KL reduction on pretraining datasets highly correlates with downstream performance.
Data selection in deep learning. Many works show the importance of data selection in the supervised or semi-supervised learning setting in vision [76; 54; 33; 36; 35; 37; 83; 85; 61; 55; 15; 71; 5] and in language finetuning [15; 54]. While most select image data from CIFAR or ImageNet, which have up to 1-10M examples, we consider selecting text data from The Pile, which has over 1.6B examples (of 128 whitespace-delimited words each). At this scale, previous methods become quite expensive since they typically require running a neural network forward pass to get embeddings [71; 76; 37], taking gradients [35; 36; 83; 61; 55], or training a reference model [54]. In contrast, we construct a simple ngram-based selection method that easily scales to internet-scale datasets. Coleman et al. [15] select data with high uncertainty under a smaller proxy neural model. They do not consider using a target dataset for estimating importance weights. However, using a neural model could be a complementary strategy for importance resampling. Other works [70; 50; 32] focus on choosing a subset that approximates training with the original dataset and require selecting data online during training. We aim to select a targeted dataset (once, before training) with different properties from the raw data (restricting the data to formal text or a specific domain). Our work also differs from active learning methods [72; 19; 84; 90; 80], which query an annotator for more labeled data. Instead, we select data for self-supervised pretraining.
Importance weighting and domain adaptation. Many methods tackle the high-dimensional importance weight estimation problem [79; 67; 12; 13]. In particular, importance weighting is classically used in domain adaptation [74; 78], where unlabeled target examples are used to adapt a model trained on labeled source data, for reweighting the loss function. However, in many modern applications the source and target are often disjoint (e.g., sketches vs. natural images), causing undefined importance weights [62; 42; 73]. We side-step high-dimensional importance weight estimation by instead working in a reduced feature space where the support of the massive web corpus should cover the target.
9 Discussion and Limitations
Feature space for importance resampling. Finding an appropriate feature space is important for DSIR. Although we find a tight correlation between downstream performance and our data metric compute using hashed n-gram features, n-grams only capture a superficial word-level overlap. Other feature extractors, such as neural models, may produce features that better capture semantics. We consider a variant of DSIR which estimates importance weights on a neural feature space in Appendix B, and find that this variant also improves by  $1 - 1.5\%$  over random selection and heuristic classification on GLUE, but our preliminary version does not improve over DSIR with hashed n-gram features. However, extracting these features is much more computationally expensive (on the order of  $D$  times more FLOPs for a  $D$ -parameter neural model), and importance weight estimation on this continuous feature space may be more difficult.
Parameterization of the importance weight estimator. In principle, both generative and discriminative approaches to estimating the importance weights should work. In a discriminative approach, regularization and calibration should be used to combat overfitting and make the predicted probabilities useful for importance resampling. We find that a generative approach requires less tuning and could also be better when the number of target examples is small, as Ng and Jordan [58] finds that Naive Bayes often performs better than logistic regression in low-sample regimes.
What is the right target distribution? When developing a domain-specific model such as Codex [11], the target dataset should be representative of the coding tasks we expect the model to be used on. However, it's unclear how exactly to collect this dataset and how much to weight each task in the target distribution. Developing better procedures for collecting the target dataset can ultimately improve the data selected by DSIR. For general-domain LMs, we follow GPT-3, the Pile, and PaLM in using formal text from Wikipedia and books as a proxy for high quality text [10; 21; 18; 14]. However, this is just a heuristic. We leave the exploration of other target distributions for general-domain LMs to future work.
Broader impacts. The impact of DSIR depends on the properties of the target data. While DSIR could amplify biases present in the target examples, with the appropriate target data, DSIR can be used to collect data that improve the training efficiency, alignment, or bias of LMs [59; 4; 40]. These benefits could reduce the environmental impact of LMs [77; 43; 46; 60] and reduce their biases and risks [9; 1; 22; 8]. For example, DSIR can be used to collect more data on underrepresented subpopulations and fine-tune the model on this data to improve model fairness.
We provide a cheap and scalable data selection framework based on importance resampling for improving the downstream performance of LMs. We also find a data metric, KL reduction, that strongly correlates with downstream performance and can provide a sanity check for data selection methods without training a model. Our work provides a step in understanding the choice of pretraining data for downstream transfer in LMs.
11 Acknowledgements
We thank Neil Band, Hong Liu, Garrett Thomas, and anonymous reviewers for their feedback. This work was supported by an Open Philanthropy Project Award and NSF IIS 2211780. SMX was supported by a NDSEG Fellowship. SS is supported by an Open Philanthropy Graduate Fellowship.
[1] Abubakar Abid, Maheen Farooqi, and James Zou. Persistent anti-muslim bias in large language models. arXiv preprint arXiv:2101.05783, 2021.
[2] Amittai Axelrod. Cynical selection of language model training data. CoRR, abs/1709.02279, 2017. URL http://arxiv.org/abs/1709.02279.
[92] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Advances in Neural Information Processing Systems (NeurIPS), 2015.
A DSIR asymptotically selects from the target
We prove that DSIR selects examples with features distributed as the target as the raw dataset size goes to infinity, assuming that the importance weights are correct up to a constant factor.
Proposition A.1. Assume that the importance weights  $w_{i}$  are proportional to the true importance weights  $\frac{p_{\text {feat }}(z_{i})}{q_{\text {feat }}(z_{i})}$ . Then as the number of raw examples  $N$  goes to infinity, the procedure returns  $k$  i.i.d. samples with features distributed according to the target feature distribution  $p_{\text {feat }}$ .
Proof. By assumption, we have importance weights  $w_{i}$  that are proportional to the true importance weights, so that  $w_{i} = C\frac{p_{\mathrm{feat}}(z_{i})}{q_{\mathrm{feat}}(z_{i})}$  for the  $i$ -th source example for some constant  $C > 0$ . First suppose that  $k = 1$ . Then,
Prob. of sampling an example with feature value  $z = \frac{\sum_{i=1}^{N} \mathbf{1}[z_i = z] w_i}{\sum_{j=1}^{N} w_j}$  (3)
$$
\begin{array}{l} = \frac {C \sum_ {i = 1} ^ {N} \mathbf {1} \left[ z _ {i} = z \right] \frac {p _ {\text {f e a t}} \left(z _ {i}\right)}{q _ {\text {f e a t}} \left(z _ {i}\right)}}{\sum_ {j = 1} ^ {N} C \frac {p _ {\text {f e a t}} \left(z _ {j}\right)}{q _ {\text {f e a t}} \left(z _ {j}\right)}} (4) \\ = \frac {\frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {1} \left[ z _ {i} = z \right] \frac {p _ {\text {f e a t}} \left(z _ {i}\right)}{q _ {\text {f e a t}} \left(z _ {i}\right)}}{\frac {1}{N} \sum_ {j = 1} ^ {N} \frac {p _ {\text {f e a t}} \left(z _ {j}\right)}{q _ {\text {f e a t}} \left(z _ {j}\right)}}. (5) \\ \end{array}
$$
For  $k \geq 1$ , we can similarly compute the probability of sampling the  $m$ -th example ( $m \in \{1, \dots, k\}$ ) as:
Prob. of sampling  $m$ -th example with feature value  $z = \frac{\frac{1}{N - m + 1}\sum_{i = 1}^{N - m + 1}\mathbf{1}[z_i = z]\frac{p_{\mathrm{feat}}(z_i)}{q_{\mathrm{feat}}(z_i)}}{\frac{1}{N - m + 1}\sum_{j = 1}^{N - m + 1}\frac{p_{\mathrm{feat}}(z_j)}{q_{\mathrm{feat}}(z_j)}}$  (6)
where for notational convenience, we re-index the raw examples after selecting each example.
For each  $m \in \{1, \dots, k\}$ , the numerator converges to  $p_{\mathrm{feat}}(z)$  as  $N \to \infty$ :
$$
\frac {1}{N - m + 1} \sum_ {i = 1} ^ {N - m + 1} \mathbf {1} [ z _ {i} = z ] \frac {p _ {\text {f e a t}} (z _ {i})}{q _ {\text {f e a t}} (z _ {i})} = \frac {1}{N - m + 1} \sum_ {i = 1} ^ {N - m + 1} \mathbf {1} [ z _ {i} = z ] \frac {p _ {\text {f e a t}} (z)}{q _ {\text {f e a t}} (z)} \rightarrow q _ {\text {f e a t}} (z) \frac {p _ {\text {f e a t}} (z)}{q _ {\text {f e a t}} (z)} = p _ {\text {f e a t}} (z) \tag {7}
$$
since  $z_{j}$  (raw features) are sampled from  $q_{\mathrm{feat}}$  (raw feature distribution). For the same reason, the denominator converges to 1:
$$
\frac {1}{N - m + 1} \sum_ {j = 1} ^ {N - m + 1} \frac {p _ {\text {f e a t}} (z _ {j})}{q _ {\text {f e a t}} (z _ {j})} \rightarrow \mathbb {E} _ {q _ {\text {f e a t}}} \left[ \frac {p _ {\text {f e a t}} (z _ {j})}{q _ {\text {f e a t}} (z _ {j})} \right] = 1. \tag {8}
$$
Therefore the features of the  $m$ -th example is sampled from  $p_{\mathrm{feat}}$  for all  $m \in \{1, \dots, k\}$ .
Intuition from a simple example. DSIR uses importance resampling to better balance the tradeoff between relevance and diversity as the samples converge to true samples from the target distribution (Proposition A.1), while no such guarantee holds for top-k selection. For intuition using a simple example, consider a raw dataset of  $n$  coin flips from a biased coin, with  $0.9n$  heads and  $0.1n$  tails. We want to filter the raw dataset to have  $k = 10$  flips from a fair coin (the target distribution). The importance weights are  $\frac{1}{2\cdot 0.9}$  for the heads examples and  $\frac{1}{2\cdot 0.1}$  for the tails examples (the tails have higher weight). If we select the top  $k$  flips according to the importance weight, we will select 10 tails, still resulting in a biased dataset. However, importance resampling balances this out, resulting in a fair dataset in expectation as  $n$  goes to infinity. We ran a simulation of the simple example with  $k = 10$  and varying raw data sizes  $n$  to see how fast the resampled dataset converges to a fair dataset with the raw data size. For raw data sizes  $n \in \{100, 200, 500\}$ , DSIR selects a dataset with  $(44\%, 47\%, 50\%)$  heads respectively, averaged over 1000 trials. Thus, DSIR converges quickly to the desired target distribution. In all cases, top-  $k$  selects a dataset with all tails.
Table 4: Accuracies on the GLUE [82] dev set for a BERT-style masked language model [17] trained on data selected from The Pile [21]. Following RoBERTa [48], for RTE, STS, and MRPC we fine-tune starting from the MNLI model instead of from scratch. DSIR outperforms heuristic classification (used by GPT-3 and PaLM) and random selection by over  $2\%$  on average. All results are averaged over 5 seeds and standard deviations are in subscripts.
MNLI	QNLI	QQP	RTE	SST-2	MRPC	CoLA	STS-B	Avg
Random selection	82.630.41	86.900.28	89.570.30	67.371.69	90.050.41	87.401.08	49.413.67	88.630.11	80.25
Heuristic classification	82.690.17	85.950.79	89.770.32	68.591.75	88.940.98	86.030.93	48.173.19	88.620.22	79.85
Top-k Heuristic classification	83.340.22	88.620.24	89.890.19	70.040.99	91.150.76	86.371.00	53.023.56	89.300.11	81.47
DSIR	83.070.29	89.110.14	89.800.37	75.092.76	90.480.57	87.700.68	54.001.34	89.170.13	82.30
Top-k DSIR	83.390.06	88.630.38	89.940.17	72.491.29	91.010.79	86.181.12	49.901.10	89.520.21	81.38
DSIR + Neural features	83.440.16	88.200.35	89.810.35	70.682.70	90.501.07	87.550.99	52.581.67	88.400.12	81.40
Figure 5: Distribution of Pile data sources for datasets selected by Left: Random selection Middle: Heuristic classification and Right: DSIR. Heuristic classification and DSIR were restricted to select only  $4\%$  of its dataset from Wikipedia, Books3, BookCorpus2, and Gutenberg.
B DSIR with a neural importance weight estimator
As a preliminary study, we test an instantiation of DSIR with an importance weight estimator based on neural features. For each example, we extract embeddings from a SentenceTransformer [66] (a11-MiniLM-L6-v2) with dimension 384. We fit the generative models for the source and target feature space with 1000- and 50-component Gaussian mixture models respectively, with diagonal covariance structure. We use this to select pretraining data for training general-domain LMs from scratch. Table 4 shows the results using this DSIR variant in the last row. On average, DSIR with neural features improves by  $1 - 1.5\%+$  over random selection and heuristic classification and is on par with top- $k$  heuristic classification and top- $k$  DSIR, but still underperforms DSIR with n-gram features. However, we believe that many aspects of this preliminary pipeline could be improved or redesigned, and that using a neural model in the importance weight estimator is a promising direction.
C Distribution of data sources for general-domain training
Figure 5 shows the distribution of data sources (ArXiv, GitHub, News, etc.) from The Pile that were selected by random selection, heuristic classification, and DSIR. Heuristic classification and DSIR aim to select formal text that are similar to text from Wikipedia or books. Note that we restricted heuristic classification and DSIR to select from data sources outside of Wikipedia and books sources (Books3, BookCorpus2, Gutenberg) for  $96\%$  of the dataset, while  $2\%$  is randomly selected from Wikipedia and the remaining  $2\%$  are selected from the 3 book sources. DSIR seems to focus mostly selecting formal text from web data such as Pile-CC (which can still be quite varied), while the other methods select from a variety of sources.
D Distribution of data sources for continued pretraining
Figure 6 shows the distribution of Pile data sources selected by DSIR for different target distributions. Each of the 4 columns represents a domain: CS papers, Biomedical text, News, and Reviews. The
Figure 6: Distribution of Pile data sources selected by DSIR for different target distributions. The four columns from left to right represent 4 domains: CS papers, Biomedical text, News, and Reviews.
Figure 7: Plot of KL reduction against average downstream F1 score of DSIR for the 8 continued pretraining datasets, where each point represents a different pretraining dataset (selected for a particular target). Pretraining data selected for CS-related target distributions tend to transfer well to datasets in other domains, while pretraining data selected for reviews transfers poorly.