INGENIOUS : Using Informative Data Subsets for Efficient Pre-Training of Language Models (2023)

Abstract: A salient characteristic of pre-trained language models (PTLMs) is a remarkable improvement in their generalization capability and emergence of new capabilities with increasing model capacity and pre-training dataset size. Consequently, we are witnessing the development of enormous models pushing the state-of-the-art. It is, however, imperative to realize that this inevitably leads to prohibitively long training times, extortionate computing costs, and a detrimental environmental impact. Significant efforts are underway to make PTLM training more efficient through innovations in model architectures, training pipelines, and loss function design, with scant attention being paid to optimizing the utility of training data. The key question that we ask is whether it is possible to train PTLMs by employing only highly informative subsets of the training data while maintaining downstream performance? Building upon the recent progress in informative data subset selection, we show how we can employ submodular optimization to select highly representative subsets of the training corpora and demonstrate that the proposed framework can be applied to efficiently train multiple PTLMs (BERT, BioBERT, GPT-2) using only a fraction of data. Further, we perform a rigorous empirical evaluation to show that the resulting models achieve up to ~99% of the performance of the fully-trained models. We made our framework publicly available at https://github.com/Efficient-AI/ingenious .

## introduction pre-trained language models (ptlms) @xcite @xcite @xcite have revolutionized the field of natural language processing (nlp), becoming the default choice for a wide array of nlp tasks. the versatility of ptlms, however, is accompanied by significant costs. for instance, it costs an estimated $12 million to train @xcite with roughly 1.2 million pounds of co 2 emissions @xcite . megatron-turing nlg @xcite is a 530 billion parameter ptlm, which is thrice the size of gpt-3 and is trained on 4480 a100 gpus and yields close to 1% performance improvements over gpt-3. by continually increasing the size of ptlms and pre-training corpora to improve generalization ability, significant additional resources and energy are consumed, resulting in dire environmental consequences @xcite . further, such large-scale resource utilization and the costs associated with ptlms create an uneven playing field for small organizations and universities, which operate with significant resource constraints. hence, a crucial step towards developing responsible, fair, and greenai @xcite involves minimizing inefficiencies and costs of training these models. significant efforts toward improving the efficiency of ptlms have ventured in directions such as optimizing the model architecture @xcite @xcite , modifications to the training pipeline @xcite and task (schick and schütze, 2021), sample efficient masking techniques for improved convergence @xcite and leveraging contextual knowledge to reduce model size @xcite . in this work, driven by the observation that the scale of the pre-training corpus contributes significantly to the training costs of ptlms, we explore the feasibility of training ptlms using highly informative subsets of the corpus. recent studies have demonstrated the feasibility of informative data subset selection for efficient deep model training for images @xcite @xcite in both supervised and semi-supervised settings. in light of this, the key question we attempt to answer is: can we efficiently pre-train language models using highly informative subsets of the training corpus without compromising performance? the first step in answering the above question is identifying informative (or representative) subsets of the underlying training corpus such that they maximize the representation of the remaining samples in the corpus. intuitively, given a set of sentences, the subsequent addition of sentences similar to existing sentences in the set yields diminishing returns. more information gains can be achieved by adding diverse, dissimilar sentences. while the classical subset selection problem is np-hard, we can leverage the diminishing gains property of submodular functions @xcite and frame subset selection as a submodular maximization problem. several recent works @xcite @xcite @xcite have formulated the subset selection problem as that of maximizing a submodular objective. however, applying existing subset selection frameworks to ptlms is nontrivial given the scale of corpora typically used for pre-training (e.g., wikipedia and common crawl consisting of hundreds of millions of sequences and billions of tokens). most of the existing methods rely on per-sample gradients, which are expensive to compute, and to the best of our knowledge, none of the previous works have considered subset selection for such large datasets. ## related work knowledge distillation and pruning based methods @xcite @xcite pre-train a smaller variant of ptlms (such as bert) with lesser capacity using the full model as teacher network. even though lighter versions such as distilbert @xcite retain ≈ 97% of the performance with up to 60% faster inference, the ptlm still needs to be completely pre-trained initially to be able to distill the lighter version. thus, the efficiency gains are restricted only to the fine-tuning and inference. other methods prune the architecture through forcing the weights with lesser magnitude to zero value during pre-training @xcite as well as during finetuning @xcite . model architecture and training task optimizations: schick and schütze (2021) have shown that smaller ptlms can achieve better performance by formulating the task input in cloze style. @xcite proposed to optimize bert pretraining through multiple optimizations related to data, model size, and optimizer choice. @xcite proposed a staged training mechanism where they start with training a relatively smaller model, which is then used for initializing the full capacity model at a later stage. @xcite identify relevant samples from the pretraining corpus based on their similarity with the task-specific dataset to train task-specific ptlm followed by fine-tuning, thus inherently suffering from the limitation of pre-training separate models for every downstream task. curriculum learning based methods employ the sequence length of training samples as a proxy for hardness.typically, shorter (easier) sequences are presented in the initial stages of pre-training followed by longer (harder) sequences at later stages @xcite . however, such methods have been shown to perform well only in limited configurations with respect to the choice of language models, stage of pre-training, etc.. ## the ingenious framework we now present ingenious -an informative data subset selection framework for pre-training language models. we summarize the training pipeline in figure 2 . we first describe the nota-tion to formulate the problem, followed by details of different steps involved in the framework. ## experiments and results we use bert @xcite , gpt-2 (radford et al., 2019) and a domain-specific version of bert -biobert @xcite as the underlying lms. specifically, we use bert-base(110m) and gpt2-small(124m). for bert, we use english wikipedia in conjunction with bookscorpus as the pre-training corpora and employ mlm and nsp tasks for pre-training following details in the work of @xcite . we perform pre-training using a batch size of 1024 for 1,000,000 steps in the case of vanilla-bert. we perform ablations over data subset sizes and number of pre-training steps for ingenious enabled pre-training and find a subset size of 25% (appendix j) with 250,000 pre-training steps (25%) as an optimal choice. we set the value of r to 25000 steps. we refer the reader to appendix g for further implementation details. for ingenious enabled pre-training of biobert and gpt-2, we discuss the implementation details and experimental results in sections 4.4 and 4.5, respectively. ## conclusions we presented ingenious, a framework for efficient pre-training of language models using highly informative data subsets, and presented a submodular optimization based algorithm. we described how it can be scaled for language models and showed its effectiveness using rigorous empirical evaluation. our future work will explore exploiting external knowledge bases to identify and reduce redundancies in the corpus and to study multi-modal training where redundant information can be spread across different modalities.