Critical Learning Periods: Leveraging Early Training Dynamics for Efficient Data Pruning (2024)

Abstract: Neural Machine Translation models are extremely data and compute-hungry. However, not all datapoints contribute equally to model training and generalization. Data pruning to remove the low-value data points has the benefit of drastically reducing the compute budget without significantdrop in model performance. In this paper, we propose a new data pruning technique: CheckpointsAcross Time (CAT ), that leverages early model training dynamics to identify the most relevantdata points for model performance. We benchmark CAT against several data pruning techniquesincluding COMET-QE, LASER and LaBSE. We find that CAT outperforms the benchmarks onIndo-European languages on multiple test sets. When applied to English-German, English-Frenchand English-Swahili translation tasks, CAT achieves comparable performance to using the fulldataset, while pruning up to 50% of training data. We inspect the data points that CAT selectsand find that it tends to favour longer sentences and sentences with unique or rare words.

## introduction in recent years, there has been significant improvement in the quality of neural machine translation (nmt) models @xcite @xcite @xcite @xcite @xcite . the success of these models is primarily due to many factors including pretraining data size, compute abundance and the ever-increasing model size. despite the gains due to growing datasets, large dataset sizes have also posed significant hurdles for maintaining data quality. significant portions of web-scraped data used for language model pretraining have been shown to be of low quality, machinegenerated spam, pornographic content @xcite . given the ever-growing size of parallel corpora, it often becomes laborious for humans to assess the quality at scale @xcite @xcite . quality of pretraining data for large nmt models often faces amplified quality issues because of extensive bitext mining approaches to automatically extracting translation pairs from either monolingual corpora @xcite @xcite @xcite or document-aligned corpora @xcite @xcite . bitext mining approaches, although far more costeffective than human curation, have been shown to introduce significant noise into datasets by relying on sub-optimal sentence embedding models @xcite @xcite to match sub-optimal translation pairs @xcite . recent work has shown that quality matters. large neural models trained on high-quality data outperform models trained on noisy data @xcite @xcite or match performance using far fewer data points @xcite @xcite @xcite . in this work, we ask if we can arrive at a rigorous estimator of data quality through data pruning. our goal is to focus on metrics which are scalable across large datasets, and which identify a subset that preserves performance. we leverage distinctive periods of training to indicate examples most critical to model generalization. specifically, we propose checkpoints across time (cat): a metric that leverages the variability in perplexity across early checkpoints to identify such critical data points when training mt models. we use this to rank data points, selecting the instances with the highest variability and continuing to train the model solely on this subset. this scoring is used to remove the data points estimated to be least important, thus creating gains in efficiency for training. cat is motivated by fundamental work in machine learning which has shown that there are distinctive periods to learning in deep neural network training. prior work shows that more common, easier features are learned earlier in training, with the most challenging features learned in the last stages of training @xcite @xcite . prior work has also shown that variance in gradients with respect to inputs @xcite or gradient patterns themselves vary depending on the period of learning @xcite . however, gradients remain expensive to leverage as a learning signal to identify relevant instances. furthermore, almost the entirety of this work has focused on a computer vision setting, with the exception of @xcite which uses differences in the model's confidence in the true class, and the variability of this confidence across epoch to cluster data. herewe explore applying the insights that there are distinctive periods of training to an nlp setting and leveraging for the purpose of pruning. we propose a more efficient method to compute signal relative to prior methods which have required computing gradients instead. our hypothesis is that those with the largest variability indicate the easy examples which show the quickest learning and gains in certainty early in training. compared to existing data pruning strategies in nmt, e.g. the popular (dual) cross-entropy method @xcite and contrastive data selection @xcite , our approach doesn't require an initial clean dataset to train on as this is one limitation that often prevents the adoption of such methods to low-resourced languages. additionally, it does not require significant time to run as we rely only on two early checkpoints rather than a training run over multiple epochs. this makes cat methods suitable for low-compute and data-scarce environments which are usually co-occurring @xcite . additionally, it eliminates dependency on other models that may not capture information from various languages and domains. we study the effectiveness of cat across diverse linguistic contexts. specifically, we utilize english-german and english-french pairs from the indo-european family using wmt datasets, and the english-swahili pair from the bantu family with automatically aligned datasets. through a series of extensive experiments, we show that cat results in significant improvements in translation quality over random pruning (randomly selecting x% subset of the training data) and other embedding-based quality estimation methods such as labse @xcite . we summarize our key contributions below: using our best cat-based method on english-german, english-french and english-swahili translation tasks, we are able to achieve comparable performance while pruning up to 50% of the training data. ## cat: checkpoints across time our goal is to extract the most valuable data points that would contribute to model generalization and train the model solely on this. our approach, while simple, leverages the difference in learning stages in deep neural network optimization. firstly, we train an nmt model on the full dataset for only a few epochs, ideally a small fraction of the total time needed for the model to converge. subsequently, we use the checkpoints from this initial stage to compute the perplexity of each data point. these perplexity scores serve as the basis for selecting candidate data points for subsequent pruning using different perplexity profiles. this approach draws upon the findings of @xcite that categorize training data relevance based on the learning dynamics during training, and @xcite which leverage the difference in gradients across training to identify easy versus challenging examples. unlike @xcite which uses computationally expensive variance of gradients, we use perplexity. for cat we can utilize the perplexity scores in two ways: we then rank all data points based on differences in perplexity between the two epochs that form our checkpoints: from the ranked data points, we select those with the highest differences between early epochs, foot_0 according to a threshold at a chosen percentile. for example, for pruning of 50% of data points, we keep all data points that rank higher than the 50th percentile. ## experimental setup to fully test the cat methods for pruning we explore the following variations: (1) different target languages, (2) different pruning levels, (3) different test domains, and (4) configurations such as model sizes and choice of checkpoints. we compare the effectiveness of the pruning against embeddingbased, llm-based, and random selection. ## results and discussion we conducted experiments on english-german and english-swahili datasets, employing various pruning techniques. we pruned 90%, 70% and 50% of the training data and compared this to training on the full set. to evaluate these models, we used wmt18 and flores test sets for german, and flores and mafand-mt as the test sets for swahili. furthermore, we applied the most effective pruning techniques to an english-french dataset to verify the consistency of performance across different language pairs. ## analysis of sentence selection for german, french, and swahili to better understand the performance gap between pruning techniques for german, french(indo-european languages) and swahili(a bantu language), we analyze the sentences selected randomly as well as those selected using cat-diff(1,5) and labse. figure 5 illustrates that longer sentences correlate with higher bleu scores for german and french when considering target sentence length. this observation of longer sequence lengths giving better scores is also seen in @xcite . most of the points in the top right quadrant are chosen by cat-diff(1,5), indicating a preference for longer sentences and performs better. however, we observe that picking the longest sentences for swahili doesn't lead to better bleu scores. we, therefore also look at the lexical diversity to distinguish the sentences selected. figure 6 shows that cat-diff(1,5) selects sentences that have more unique and rare words as opposed to german which is less lexically diverse and thus results in lower mean frequency. similarly, cat-diff(1,5) also selects more unique and rare words for swahili. although labse performs relatively better in selecting more unique and rare words for french, it still falls short compared selection of words using cat-diff(1,5) for german. this implies that cat-diff(1,5) selection is far more complex than selecting the longest sentences. we leave a deeper investigation of sentences selected using various techniques for future studies. ## conclusion identifying optimal subsets of a given dataset for training a machine translation model is a crucial problem that enables data efficiency and saves on compute. in this paper, we present checkpoints across time (cat) methods as efficient techniques to prune datasets. cat uses the change in perplexity scores of individual examples over the first few training epochs to find good subsets. using 3.8m sentences for english-german, english-french and english-swahili, we show that cat techniques offer highly computationally efficient pruning that achieves above 75% of full performance even when 90% pruning is applied to all target languages. we also show that labse, when used to select training examples from an automatically aligned dataset provides strong signals on examples crucial for model generalization as the examples selected perform on par or surpass training on the full dataset. interestingly, we find that utilizing llms for pruning does not offer benefits in selecting data for german but does for swahili.