Scalable Data Ablation Approximations for Language Models through Modular Training and Merging (2024)

Abstract: Training data compositions for Large Language Models (LLMs) can significantly affect their downstream performance. However, a thorough data ablation study exploring large sets of candidate data mixtures is typically prohibitively expensive since the full effect is seen only after training the models; this can lead practitioners to settle for sub-optimal data mixtures. We propose an efficient method for approximating data ablations which trains individual models on subsets of a training corpus and reuses them across evaluations of combinations of subsets.In continued pre-training experiments, we find that, given an arbitrary evaluation set, the perplexity score of a single model trained on a candidate set of data is strongly correlated with perplexity scores of parameter averages of models trained on distinct partitions of that data. From this finding, we posit that researchers and practitioners can conduct inexpensive simulations of data ablations by maintaining a pool of models that were each trained on partitions of a large training corpus, and assessing candidate data mixtures by evaluating parameter averages of combinations of these models. This approach allows for substantial improvements in amortized training efficiency – scaling only linearly with respect to new data – by enabling reuse of previous training computation, opening new avenues for improving model performance through rigorous, incremental data assessment and mixing.

## introduction as large language models (llms) and their training corpora have grown in scale, it is increasingly costly not only to train an llm given a fixed recipe, but also to develop recipes for training new language models with improved capabilities. design decisions can span many factors such as model architecture, optimization techniques, and others; one critical aspect of llm development is training data composition. ## modular training for data ablations we characterize our proposed approach and its asymptotic complexity vs the traditional method for conducting data ablations. the asymptotic efficiency advantage of our method is comparable to that found comparing a naive recursive implementation of an algorithm with combinatorial complexity, to one that uses memoization. our method analogously benefits from an effective decomposition of a large training corpus into subsets whose corresponding trained models are reused repeatedly across many candidate data mixtures ("overlapping subproblems"). our method also similarly requires additional cache storage for component models, relatively inexpensive compared to additional compute. ## data the training datasets for our main experiments are derived from the semantic scholar open research corpus (s2orc) @xcite and m2d2 wikipedia @xcite . additionally, for a given model size, we begin by pre-training a model on a general "seed" corpus of gutenberg books and english wikipedia. all models of the same size used in the current study are initialized from a copy of this seed model before being trained on a combination s2orc and m2d2 wikipedia text. we evaluate our models using held-out sets from the continued pre-training data, as well as a subset of the validation sets from paloma @xcite , an evaluation dataset for assessing language model fit on various domain-specific text. as ours is an exploratory study, we experiment primarily at 1) smaller scales and 2) with data mixtures of carefully constructed domains within datasets of continual pre-training scale. data is partitioned topically and temporally in the case of s2orc, and topically in m2d2 wikipedia foot_3 . for m2d2 wikipedia, we use the l1 (top-level) domains as originally defined by @xcite with a median of 282 million tokens each. as s2orc is not pre-partitioned, we construct 22 partitions p for s2orc based on metadata for field of study (fos), and 128 "base unit" partitions p ′ of roughly equal size, around 287 million tokens each, based on fos and year. some documents are tagged with multiple fields of study (see all seed and continued pre-training data was decontaminated against paloma evaluation data, and all continued pre-training data was decontaminated against its respective evaluation splits. following @xcite 's standard for decontamination on paloma @xcite , a bloom filter was used to exclude all documents containing least one "contaminated" paragraph (at least 13 tokens long and found in the evaluation set). in general, we use s2orc and m2d2 wikipedia as examples of datasets at continued pre-training scale that can be partitioned along natural inherent groupings in their distributions, but we encourage future work exploring different data. ## methods for our base experiments, we train small (130m parameter) decoder-only models with a palm-like architecture, on independent partitions of the semantic scholar open research corpus (s2orc) @xcite and m2d2 wikipedia @xcite , defined as discussed in §3. in a subset of our experiments, we involve larger (1.1b parameter) models. models share an initial optimization trajectory of 42 billion tokens of "seed" pre-training on a general corpus of 10 billion tokens consisting of gutenberg books and english wikipedia. each model is then trained on a specific partition within s2orc or m2d2 wikipedia for some multiple of 1 billion tokens, depending on the experimental setup. we evaluate models using per-token perplexity scores on domain-specific validation sets as well as out-of-domain sets from paloma @xcite ; all perplexity scores reported in this study are on held-out sets. we use the gpt-neox-20b tokenizer @xcite as it is used by @xcite with the addition of 3 special tokens for masking personally identifiable information 6 . see appendix a.3 for more detailed reporting on training settings, hyperparameters, and hardware. ## experiments in our experiments, we have: partitions p ′ ∈ c are disjoint, but we have no evidence to suggest that this is a strict requirement. 7 total number of partitions is @xmath0, except in §5.1 where we sample partitions directly from {p ′ ∈ c} and so @xmath1. the set of evaluation domains in each experiment includes 1) in-domain datasets from each data mixture ⃗ d's respective combined held-out dataset and 2) nine domains from paloma as ood or domain-shifted data. ## related work pretraining mixtures existing literature includes many open corpora that are presented as viable training data mixtures (or components of them) for pre-training llms @xcite @xcite @xcite . the individual components of these mixtures and their relative sizes are typically chosen based on some intrinsic measure of data quality as it is often prohibitively expensive to perform thorough data ablations to create mixtures optimizing for downstream performance. efficient data selection given that data ablations on large language models is expensive, one class of approaches relies on approximating them on smaller models. relevant work studies scaling laws for model parameters vs training tokens @xcite , empirical effects of including or excluding different sources of data @xcite , and the effects of training over multiple epochs vs new training tokens @xcite . previous work has also explored improving domain-specific fit via continued pre-training @xcite , predicting domain fit using lexical features @xcite , or improving general test-time adaptation via dynamic data selection, either by distributionally robust optimization with a small proxy model @xcite or online using a multi-armed bandit approach @xcite . additional previous works aim to adapt to known downstream tasks via data selection, including at the individual example level @xcite or even by explicitly fine-tuning models on many tasks @xcite . notable recent and concurrent work proposes scalable influence functions, traditionally used at only very small scales, as a method for selecting better training data @xcite . another concurrent preliminary work @xcite proposes a training-free method of selecting data to improve performance on downstream tasks. see @xcite for a survey on data selection methods. we note that our strategies may be compatible with many of the existing methods for dataset selection, potentially leading to cumulative efficiency improvements -in particular, we show that our proxy metrics are compatible with @xcite 's strategy of predicting larger models' performance on a data mixture using a smaller, proxy model's performance. model averaging merging models via weightspace averaging is more commonly done in the finetuning stages of language model training @xcite , typically for improving robustness and mitigating cross-task interference. when used for pretraining language models @xcite , the goal is to efficiently adapt to new domains at inference time. to the best of our knowledge, ours is the only work that studies model merging for efficient data ablations. linear mode connectivity relevant to our work, the mode connectivity perspective argues models share behavior when linearly connected on the loss surface @xcite @xcite . notably, @xcite previously demonstrated that model ensembling is most effective when combining models on the equivariant loss curve on the loss surface. this applies to our work as all our models share a seed training phase, and we can subsequently presume each expert is close on the loss surface. ## conclusion we lay groundwork for a promising new strategy whereby one can efficiently simulate principled, fine-grained ablations of training datasets and data mixtures. we are hopeful that researchers and practitioners can follow our recommendations to find beneficial data mixtures for their own models and data. moreover, we encourage others to experiment in additional settings and report observations specific to their own assumptions and settings, particularly in relation to larger scales of models and data and limitations of the current method.