B aby L lama-2: Ensemble-Distilled Models Consistently Outperform Teachers With Limited Data (2024)

Abstract: We present BabyLlama-2, a 345 million parameter model distillation-pretrained from two teachers on a 10 million word corpus for the BabyLM competition. On the BLiMP and SuperGLUE benchmarks, BabyLlama-2 outperforms baselines trained on both 10 and 100 million word datasets with the same data mix, as well as its teacher models. Through an extensive hyperparameter sweep, we demonstrate that the advantages of distillation cannot be attributed to suboptimal hyperparameter selection of the teachers. Our findings underscore the need for further investigation into distillation techniques, particularly in data-limited settings.

## introduction with frontier model training runs using beyond 10 25 flops @xcite , training efficiency has become a billion-dollar question. humans are vastly more sample efficient than current large language models (llms). for example, a typical 13-year-old child has been exposed to less than 100 million words (extrapolating from @xcite ), whereas llama-3.1 has been trained on 15.6 trillion text tokens. the goal of the babylm challenge @xcite is to optimize pretraining given dataset limitations inspired by human development. in this work, we present our contribution to the babylm challenge (strict-small track), with the following main results: the success of distillation pretraining, i.e. pretraining from scratch with distillation loss, in our experiments highlights its potential as a powerful technique for improving model performance, especially in data-limited settings. while our findings are promising, they also raise intriguing questions about the nature of knowledge distillation and its interaction with pretraining objectives. further investigation into these areas could yield valuable insights for the development of more sample-efficient language models. ## related work the first edition of the babylm challenge, which aims to optimize language model pretraining under data constraints inspired by human language acquisition, prompted numerous works on sampleefficient pretraining. for a detailed summary of all contributions, see the review by @xcite . outside the babylm context, relatively few works address training on limited language datasets. notable exceptions include @xcite , who studied the scaling of data-constrained llms. their main finding is that training for more than 4 epochs leads to diminishing returns. @xcite trained fingpt on more than 30b tokens in finnish language. although resourceconstrained, this dataset is significantly larger than that of the babylm challenge. a sample-efficient modification of bert architecture was proposed by @xcite , with a model trained on a 100m word dataset from the british national corpus outperforming the original bert model. the existing literature on training small models often focuses on models deployable on edge devices, such as mobilellm @xcite . however, these works typically concentrate on deployment efficiency rather than sample efficiency. knowledge distillation has recently attracted significant attention, primarily for deployment efficiency reasons (see @xcite for a systematic review). typically, this involves using large frontier models as teachers to train smaller student models. in contrast, babyllama-2 utilizes distillation for sample-efficient pretraining, using similar-sized teacher models trained on the same limited dataset. a similar phenomenon, where a student model outperforms its teachers when distilled from models with identical architecture and trained on the same dataset, was observed in "born-again neural networks" @xcite . however, this work did not focus on the data-limited regime and it used lstm variants (instead of transformers) for language modeling. ## background knowledge distillation, introduced by @xcite , is a technique for transferring knowledge from a "teacher" model to a "student" model. the core idea is to train the student to mimic the logit distribution (soft targets) produced by the teacher, rather than just the hard labels of the training data. the distillation loss combines the standard crossentropy loss with the soft target loss: ## model architecture. previous experiments have shown that the llama architecture @xcite , featuring rope and a swiglu non-linearity, requires fewer epochs to reach minimal loss compared to gpt-2 or gpt-j architectures @xcite . after training a family of llama models ranging from 16m to 728m parameters, we converged on a specific 345m model architecture suggested in mobilellm @xcite pretraining approach. the particularity of the babyllama-2 model is to be distilled from an ensemble of teacher models, using the distillation loss (1). the teacher models share the same architecture and are pretrained on the same dataset using the standard cross-entropy loss. the student model is then pretrained with the same hyperparameters, using the mean teacher logits zt in the distillation loss l distill (y, z s , zt ). ## experimental setup dataset. we use the 10 million word babylm-2 dataset @xcite , that we split into 9.5m train and 0.5m validation splits, as well as the accompanying 10m word "dev" dataset, that we use as a test split. while the validation split is used to perform the hyperparameter optimization, 2 the test split is used solely for the purpose of reporting the final cross-entropy loss. each dataset is composed of six files, corresponding each to a different type of (english) language that a child is likely to be exposed to, such as transcribed childdirected speech, children's books, subtitles, or simple wikipedia. the relative fractions of these files differ slightly between, on the one hand, the train and validation splits and, on the other, the test split, which is therefore slightly out of distribution. we have experimented with the fineweb-edu dataset @xcite but have observed that models trained on the babylm-2 dataset reach better blimp scores (see appendix c for more details). training. the teacher models are pretrained using the trainer class from the huggingface transformers library, using the hyperparameters listed in all models share the same byte-pair encoding (bpe) tokenizer with a vocabulary size of 16000 trained on the training split of babylm-2 dataset. hyperparameter value learning rate 7 • 10 -4 number of epochs 8 batch size 128 weight decay 5 distillation t 1 distillation α 0.5 hyperparameter sweep. to exclude the possibility that the student model babyllama-2 outperforms its teachers due to a suboptimal choice fore we would consider using the dev split from babylm-2 as a violation of the rules of the challenge. of course, it means that we trained only on 95% of the tokens, and could potentially improve our results further. in total, we trained 265 models as part of the sweep, amounting to 26 gpu-days. while the sweep produced some runs that perform noticeable better than the teachers trained with the parameters in benchmarks. we evaluate the performance of the teacher and student models on the benchmarks suggested by the organizers of the babylm challenge. those include zero-shot benchmarkssuch as blimp @xcite , which focuses on linguistic knowledge in english, and ewok @xcite , focusing on world knowledge -as well as the suite of fine-tuning benchmarks superglue @xcite ) about language understanding. for the latter, the finetuning hyperparameters are optimized using a separate sweep for each task (totalling 1293 runs and 37 gpu-days). the optimal parameters, listed in baseline models. the organizers of babylm-2 have provided two baseline models: ltg-bert @xcite , (encoder-only) and babyllama (timiryasov and tastet, 2023a) (decoder). both models were re-trained by the challenge organizers on both the 10m and 100m word datasets. ltg-bert modifies the original bert architecture by utilizing the pre-norm variant of the transformer with geglu feed-forward layers and by disentangling positional information from token embeddings. the highest performing solution of the 2023 edition of the babylm challenge, elc-bert (charpentier and samuel, 2023), is based on this architecture. on the other hand, babyllama (the highest-performing decoder model) uses the standard llama architecture @xcite , but a modified training procedure, following a similar approach to the one presented here. however, in contrast to babyllama-2, it was distilled from two larger teachers with two different architectures (gpt and llama), and had six times less parameters. since babyllama-2 aims to demonstrate the validity of the ensemble distillation method itself, it uses same-size, homogeneous models in order to remove potential confounding factors. in addition to the baseline models, we vary the number of teachers between 1 and 3, and compare babyllama-2 to the ensemble formed by the two teacher models (applying softmax to the averaged logits zt and letting the gradient flow back into both teachers during fine-tuning, with the same training hyperparameters as for babyllama-2). when evaluating the original babyllama on the superglue benchmarks, we fine-tune it again using the hyperparameters reported in @xcite , and successfully reproduce all of its scores. ## results figure 1 summarizes the performance of the models considered in section 5 with respect to the various evaluation metrics: the cross-entropy loss evaluated on the held-out test set, the blimp scores for the "filtered" and "supplement" subsets of evaluation tasks, and the mean superglue score. the ewok benchmark is not shown, since the performance of our models and of the baselines trained on 10m words is consistent with random chance, hinting that all these models have extremely limited world knowledge, if any. distributions. violin plots are used in order to quantify the variability across model initializations, with a minimum of 5 runs per model. each subplot shows a different metric, with the y-axis listing the various models considered: the teacher models, pretrained without distillation; the student models pretrained with one, two or three teachers; the direct ensembles formed by averaging the logits of two teachers; the baseline models for the 2024 babylm challenge; and the 265 models from the hyperparameter sweep. no violin is shown for baseline models, since they do not have an associated distribution. similarly, running fine-tuning benchmarks for all the models from the sweep would have been computationally prohibitive, therefore the superglue distribution associated with the sweep is not present, with only the best checkpoint being shown. ## conclusions in this study, we prioritized investigating the robustness of the distillation approach over architectural modifications or dataset curation. our findings demonstrate that a 345m parameter model, distillation-pretrained on 9.5m words, outperforms models of the same size and architecture pretrained in the usual way. we carried out a systematic analysis to exclude the possibility that the performance gains were due to a single fortunate initialization or suboptimal teacher model hyperparameters. through an extensive hyperparameter sweep and the training of multiple teacher and student models, we established that distillation-pretraining consistently yields superior performance. our results indicate that distillation-pretraining is an effective method for achieving high performance without the need for meticulous hyperparameter tuning, at least within the data-limited regime. the scalability of this approach to larger datasets and model sizes, as well as its applicability to other modalities, remains an open research question.