Data Selection for Language Models via Importance Resampling

By Sang Michael Xie, Shibani Santurkar, Tengyu Ma, Percy Liang. Neural Information Processing Systems, 2023

356 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence

View Paper: https://www.semanticscholar.org/paper/74013b7cfa0fc524803350fca51341004565eb22

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Fine-tuning Aligned Language Models Compromises Safety, Even When Users Do Not Intend To! (2023, methodology): “An alternative mitigation could be stricter pruning or selection of pre-training data (Xie et al., 2023), following the method used to reduce toxicity in pre-trained LLMs (Gehman et al., 2020).”
A Survey on Large Language Models for Code Generation (2024, unknown): “between smaller and larger LLMs. 7B/13B models outperform the 70B model.”
The MiniPile Challenge for Data-Efficient Language Models (2023, background): “We emphasize that our goal is not to attain state-of-the-art performance on GLUE/SNI; specifically for these downstream benchmarks, data selected from target distribution [60] could be better suited.”
From Generation to Judgment: Opportunities and Challenges of LLM-as-a-judge (2024, unknown): “As data selection (Xie et al., 2023; Albalak et al., 2024) becomes an increasingly popular research area for improving the efficiency of LLMs’ training and inference, it also holds the potential for enhancing LLMs’ evaluation.”
LESS: Selecting Influential Data for Targeted Instruction Tuning (2024, methodology): “In our empirical observations, we noticed that DSIR tends to prefer shorter sequences, which can harm performance (see Table 12 for similar evidence).”
DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining (2023, background): “Du et al. [13]; Xie et al. [55] show that ﬁltering the data at an example level for high-quality text that look like Wikipedia and books can signiﬁcantly improve downstream performance for LMs.”
CodeGen2: Lessons for Training LLMs on Programming and Natural Languages (2023, methodology): “Xie et al. (2023) introduces a data selection method based on importance resampling which allows to mix datasets of various sizes, however, the evaluation only covers encoder-only models.”