PaperlayerView exact revision on arXiv

Source-linked AI summary

Data Selection for Language Models via Importance Resampling

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

arXiv:2302.03169v3cs.CLcs.LG

TL;DR

Pretraining-data selection must match a large raw corpus to a desired target distribution, but text makes direct importance estimation difficult. DSIR estimates importance weights in a reduced feature space and resamples data accordingly, achieving scalable selection and strong downstream results. Its hashed n-gram KL-reduction metric also tracks downstream performance, while target-distribution choice and richer feature spaces remain limitations.

  • Problem

    Selecting pretraining data is important, but high-dimensional text makes direct importance-weight estimation difficult and existing approaches rely on heuristics or manual curation.

  • Method

    DSIR maps raw and target data into a reduced feature space, estimates importance weights there, and resamples raw examples according to those weights.

  • Results

    DSIR is comparable to expert curation for domain-specific continued pretraining, improves general-domain GLUE performance by 2–2.5% over random selection and heuristic classification, and KL reduction correlates with downstream performance at Pearson r=0.82.

  • Takeaways & Limitations

    DSIR provides a cheap, scalable alternative to manual curation, while KL reduction offers a model-free sanity check for data-selection methods.

  • Takeaways & Limitations

    The general-domain target uses Wikipedia and books as a heuristic proxy for high-quality text, and neural feature spaces are more computationally expensive than hashed n-grams.

Abstract

from arXiv · show

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.

1 Introduction

The paper frames pretraining-data selection as matching a large raw dataset to a desired target distribution, replacing heuristic or manual approaches with scalable importance resampling. DSIR estimates weights in a reduced feature space and improves downstream performance in both domain-specific and general-domain settings.

  • Pretraining-data choice is critical for language-model performance under a fixed compute budget.
  • DSIR selects a subset of a large raw dataset that matches a smaller target dataset in a chosen feature space.It estimates importance weights from raw and target data, then resamples raw examples according to those weights.
  • 100M documents from The Pile can be selected in 4.5 hours using hashed n-gram features.
  • KL reduction correlates with average downstream performance across 8 data-selection methods, including expert selection, with Pearson r=0.82.
  • Domain-specific continued pretraining: DSIR improves over RoBERTa without continued pretraining by 2% on average across 8 datasets and is comparable to expert-curated continued pretraining.
  • General-domain pretraining: DSIR improves over random selection and heuristic classification by 2–2.5% on average on GLUE for general-domain pretraining.

2 Setup

The setup formalizes data selection as choosing k examples from a large raw distribution to resemble a target distribution. It contrasts this goal with heuristic classification, which independently filters examples and may reduce diversity.

  • The task is to select k examples, with k≪N, from a raw dataset of N examples so they resemble a target distribution.
  • Heuristic classification trains a binary classifier to estimate whether each raw example comes from the target distribution.
  • Examples are independently selected when their predicted target probability exceeds a noisy Pareto-based threshold.The noise is intended to combat concentration on high-probability modes of the target distribution.
  • Because examples are kept or discarded independently, selecting a desired number of examples requires repetition or tuning the Pareto-shape parameter.

3 Data Selection with Importance Resampling

DSIR makes importance resampling tractable for text by mapping raw and target examples into a structured feature space. It learns feature distributions, computes density-ratio weights, and samples without replacement.

  • DSIR maps each input into a feature space where raw and target feature distributions can support tractable importance estimation.
  • The framework learns estimated target and raw feature distributions from held-out featurized examples.
  • For each raw example, DSIR computes an importance weight as the estimated target feature probability divided by the estimated raw feature probability.
  • It samples k raw examples without replacement with probabilities proportional to their importance weights.Sampling without replacement avoids duplicate selections and can be implemented with the Gumbel top-k trick.

4 DSIR with Hashed N-gram Features

The paper instantiates DSIR with hashed unigram and bigram counts, using simple bag-of-ngrams distributions for scalable estimation. This implementation supports large-scale selection from The Pile.

  • Hashed n-gram features: Hashed n-gram features map unigrams and bigrams into 10,000 buckets and represent each document by bucket counts.Hash collisions introduce noise, but the representation incorporates both unigram and bigram information efficiently.
  • Bag-of-ngrams model: DSIR uses bag-of-ngrams models to parameterize raw and target feature distributions with probability vectors over hash buckets.
  • Bag-of-ngrams model: The bag-of-ngrams distribution parameters are estimated by counting hashed-bucket frequencies.
  • Scalability: DSIR selects 100M documents from the full Pile dataset in 4.5 hours using 10k hashed buckets and one CPU node with 96 cores.
  • Evaluation: Table 1 compares F1 scores for continued pretraining across 8 downstream datasets from 4 domains under equal training compute.Results are averaged over 5 seeds; most datasets use macro-F1, while ChemProt and RCT use micro-F1.

5 Selecting Data for Domain-Specific Continued Pretraining

The section evaluates DSIR for domain-specific continued pretraining across eight datasets and four domains, comparing it with random selection, heuristics, manual curation, and variants. DSIR generally matches or exceeds these baselines, while target-data mismatch and feature choices affect downstream F1.

  • Main results: DSIR improves over random selection by 1.2% on average and manually curated data by 0.3%.These results support the potential for automatic selection to replace manual curation.
  • Main results: DSIR improves over heuristic classification by 0.9% and is comparable to top-k heuristic classification.Top-k heuristic classification may suit domain-specific selection when diversity is less important.
  • Main results: Random selection improves over no continued pretraining by 0.4% on average, while all targeted selection methods improve over random selection.
  • Ablations: Generative importance-weight estimators outperform discriminative estimators, while discriminative approaches still outperform random selection by 0.6%.The discriminative DSIR variant underperforms DSIR by 0.7%.
  • Ablations: DSIR with unigram features underperforms DSIR with hashed n-grams by 0.26%, although it remains comparable to manual curation.Including bigrams further improves the relevance of selected data.
  • Cross-domain transfer: Choosing the worst pretraining dataset for a downstream task reduces F1 by 6% on average and by 30% for HyperPartisan.The results are averaged over five seeds.
  • Cross-domain transfer: Within-domain F1 is 82.9%, 1.7% higher than the 81.2% cross-domain F1.The comparison excludes pairs where pretraining data was selected for the fine-tuning task.

6 KL Reduction on Hashed N-grams Predicts Downstream Performance

The section introduces KL reduction to assess whether hashed n-gram features preserve target-relevant information. KL reduction strongly correlates with downstream performance and can provide a cheap sanity check for data-selection methods.

  • KL reduction measures how much data selection reduces empirical KL divergence to target distributions compared with random selection.It averages this reduction over a set of target feature distributions.
  • In the continued-pretraining setting, the raw feature distribution is from The Pile and the target set contains feature distributions from eight downstream tasks.
  • Average downstream F1 correlates strongly with hashed n-gram KL reduction across eight selection methods, with Pearson r = 0.82.The comparison uses eight target distributions and feature-distribution estimates from 100k samples per dataset.
  • DSIR with hashed n-gram features achieves the highest KL reduction and the best average downstream F1.
  • Hashed n-gram features appear to preserve most information needed for selecting target-relevant data, and KL reduction can be computed without training a language model.The authors present KL reduction as a sanity check for future data-selection methods.

7 Selecting Data for Training General-Domain LMs

The section evaluates selecting formal text from Wikipedia and books for general-domain language-model pretraining. DSIR selects more formal text and improves GLUE performance over random selection and heuristic classification.

  • The target distribution for general-domain language models is formal text from Wikipedia and books.
  • DSIR selects data from The Pile alongside random selection and heuristic classification, with top-k variants as ablations.
  • DSIR qualitatively selects more formal text than random selection and heuristic classification.Random selection contains many code examples unlike the target, while heuristic classification appears slightly too diverse.
  • 82.3% average GLUE accuracy: DSIR improves over random selection by 2% and heuristic classification by 2.5%.Heuristic classification is 0.5% lower than random selection in this evaluation.
  • DSIR without top-k achieves a 0.8–0.9% higher average than top-k variants across datasets.All top accuracies across datasets are achieved by DSIR or top-k DSIR.

8 Related Work

Related work covers heuristic and classical data selection, broader data-selection research, and importance weighting for domain adaptation. It distinguishes DSIR's importance-resampling approach from methods that select examples by scores or discrepancies.

  • GPT-3, PaLM, and related systems use heuristic filtering or classification to select text closer to formal Wikipedia and book text.Heuristic classification trains a binary classifier on n-gram features to estimate whether inputs come from the target distribution.
  • BM25 retrieval can return many exact duplicates, creating oversampling and deduplication issues when selecting a desired number of unique documents.A preliminary test found 1.8M unique documents among 6.1M retrieved documents.
  • Classical Moore-Lewis selection uses cross-entropy differences from target and raw n-gram language models, whereas DSIR uses importance resampling to balance diversity and relevance.Top-k DSIR underperforms DSIR by 0.9% on GLUE.
  • Data-selection research also includes supervised and semi-supervised selection in vision and language fine-tuning.
  • Importance weighting is classically used in domain adaptation to reweight losses using unlabeled target examples, but modern source and target settings differ.

9 Discussion and Limitations

The discussion examines feature-space and estimator choices, target-distribution design, and broader impacts. It notes computational and specification limits while describing how target data can shape DSIR's effects.

  • Hashed n-grams capture superficial word-level overlap, while neural features may better capture semantics but cost substantially more computation.The neural-feature variant improves 1–1.5% over random selection and heuristic classification on GLUE but does not improve over hashed n-grams.
  • A generative importance-weight estimator requires less tuning and may work better when the number of target examples is small.
  • The quality of DSIR's selected data depends on how representative the target dataset is and how tasks are weighted within it.For general-domain models, Wikipedia and books are used as a heuristic proxy for high-quality text.
  • DSIR may amplify biases present in target examples, while appropriate target data could support improvements in training efficiency, alignment, or bias.The paper gives underrepresented-subpopulation data collection as an example for improving model fairness.

10 Conclusion

The conclusion presents DSIR as a cheap, scalable importance-resampling framework for improving downstream language-model performance. It also highlights KL reduction as a model-free sanity check for data selection.

  • DSIR provides a cheap and scalable data-selection framework based on importance resampling.
  • KL reduction strongly correlates with downstream performance and can sanity-check data-selection methods without training a model.

A DSIR asymptotically selects from the target

Under correct importance weights, DSIR asymptotically returns samples whose reduced features follow the target distribution. Its importance-resampling procedure balances relevance and diversity, unlike top-k selection, and supports strong downstream results.

  • Proposition 1: As N goes to infinity, DSIR returns k i.i.d. samples with features distributed according to the target feature distribution, assuming weights are proportional to true importance weights.The assumption allows weights to differ from the true ratios only by a positive constant factor.
  • Intuition from a simple example: Importance resampling balances target relevance and diversity, whereas top-k selection can oversample a target mode.In the coin-flip example, top-k selects all tails, while resampling converges toward the fair target.
  • Empirical results: DSIR outperformed heuristic classification and random selection by over 2% on average on GLUE accuracies.The comparison used a BERT-style masked language model trained on data selected from The Pile, with results averaged over five seeds.
  • Intuition from a simple example: For raw sizes n=100, 200, and 500, DSIR selected datasets averaging 44%, 47%, and 50% heads, respectively, toward the 50% fair target.These averages were computed over 1,000 trials; top-k selected all tails in every case.
  • Empirical results: A neural-feature DSIR variant improved over random selection and heuristic classification by 1–1.5%+ on average but underperformed DSIR with n-gram features.The neural estimator used SentenceTransformer embeddings and Gaussian mixture models.

D Distribution of data sources for continued pretraining

DSIR selects domain-dependent source distributions for continued pretraining, and cross-domain transfer is asymmetric. Its continued-pretraining gains on formal-text targets exceed the reported baselines on GLUE.

  • Distribution of data sources: Data-source distributions selected by DSIR are similar for target distributions within the same domain.The compared domains are CS papers, biomedical text, news, and reviews.
  • Distribution of data sources: CS-target selection produces the most diverse source distribution, while biomedical and news targets produce particularly different distributions.These patterns describe the source composition of DSIR-selected Pile data across the four domains.
  • Continued pretraining results: DSIR improves average GLUE performance by 0.4–0.7% over all baselines when the target distribution is formal text.The continued-pretraining results are averaged over five fine-tuning seeds.
  • Cross-domain transfer: KL reduction remains a strong indicator of downstream F1 across the eight continued-pretraining datasets, with CS-selected data transferring best and review-selected data hurting performance.The figure also indicates that transfer between domains is highly asymmetric.
  • Continued pretraining results: DSIR improves over random selection by 0.4% and over BERT-base by 0.6%, reaching almost 84% on the GLUE development set.Heuristic classification decreases performance relative to random selection by 0.2% on average in this setting.

F Data selection details

The experiments select and preprocess large Pile subsets using hashed n-gram methods, controlled training budgets, and standardized fine-tuning. The data-selection pipeline combines scalable importance estimation with explicit quality filtering and distribution measurement.

  • Data preprocessing: The Pile is divided into 128-word chunks, yielding 1.7B examples after reserving one validation chunk.Only the last 29 of 30 random Pile chunks are considered for selection.
  • Selection methods: Heuristic classification uses a bigram fastText classifier, while importance resampling uses lowercased unigram and bigram counts hashed into 10k buckets.The importance-resampling model treats hashed counts as a 10k-word vocabulary for a bag-of-words generative model.
  • Selection methods: Gumbel top-k samples k examples without replacement by adding independent Gumbel noise to log importance-weight ratios and selecting the highest scores.The procedure can be parallelized for score computation and uses linear-time introselect for top-k selection.
  • General-domain sampling: General-domain selection first chooses 102.4M examples, concatenates pairs into 51.2M examples, and mixes in Wikipedia and book examples to balance sources.The target is Wikipedia plus BookCorpus2, Gutenberg, and Books3; 98.4M examples come from non-Wikipedia and non-book sources.
  • Training setup: Training from scratch processes 26B tokens, whereas continued pretraining from BERT-base processes 26B additional tokens after 43B initial tokens.Each run takes about 50 hours from scratch and about 60 hours for continued pretraining.
  • Quality filtering: The quality filter keeps 52% of examples when all filters are applied simultaneously.The authors interpret length filtering as the main filtering operation and as a possible quality proxy.
Loading 2302.03169v3…