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

TL;DR

The paper tackles the prohibitive cost of exhaustive data-ablations for pretraining large language constructs by introducing modular training on base data partitions and evaluating mixtures through parameter averages. It demonstrates that perplexity on arbitrary evaluation domains can be predicted by proxy perplexities from merged base-unit models, enabling linear-scaling assessments of new data. The approach yields substantial efficiency gains and shows that proxies trained on smaller scales can inform decisions for larger models, while providing guidance on when macro- vs micro-merged proxies are most reliable. The work discusses scalability, limitations, and directions for future research, including larger-scale experiments and downstream-task evaluations, to broaden applicability and impact.

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.

Figures (11)

• Figure 1: Given a corpus containing multiple subsets of data, a traditional approach to studying the effects of training on different data mixtures calls for training models on each candidate data mixture. Our strategy reuses training on data shared across candidate mixtures, by 1) conducting modular training of many models on equally sized "base unit" data partitions and 2) performing evaluation on parameter averages of model combinations. We find that this yields useful proxy metrics for predicting perplexity scores on arbitrary evaluation domains, such that we can simulate comprehensive data ablation studies at a fraction of the cost.
• Figure 2: Our method (green) has only linear runtime complexity with respect to the number of unique partitions in a corpus. This allows us to simulate comprehensive data ablations at extremely low cost compared to naive training on all possible partition combinations.
• Figure 3: seq(uentially trained), merged, and ind(ividual) models trained on random pairs of S2ORC base partitions (left), pairs of M2D2 Wikipedia partitions (middle) and triples of S2ORC base partitions (right), evaluated on respective held-out sets of the same (top) and on nine subsets of Paloma (bottom). Each point lies along the same $y-$value for seq score as the other proxy evaluation score(s) for the same data mixture. For held-out sets of the training data, the seq performance on the combined dataset correlates most strongly with the average of the component models' respective in-domain evaluations ("ind ID"), compared to merged model performance or mean ind scores. However, on OOD Paloma data, the Pearson's correlation is highest between seq and merged models. See Table \ref{['tab:pearson_abridged']} for correlation values and §\ref{['subsec:figure_info']} for definitions of terms.
• Figure 4: seq vs. proxy macro-averaged Paloma perplexity scores of models trained on differently sized partitions $\mathcal{P}$. We see that the macro-merged scores are the most reliable metric (Pearson's correlation 0.984).
• Figure 5: seq vs. proxy macro-averaged Paloma perplexity scores of models trained on data mixtures containing data from multiple sources. As in §\ref{['subsec:uneven_parts']}, it is beneficial to merge models trained on similar amounts of data, but here, micro-merged models are more useful -- 0.989 corr. with seq, vs. 0.963 (merge of the seq models for each $\mathcal{P}$), 0.916 (mean ind scores), 0.807 (macro-merged), or 0.759 (mean merged scores).