B aby LM Challenge: Experimenting with Self-Distillation and Reverse-Distillation for Language Model Pre-Training on Constrained Datasets (2024)

Abstract: Language models (LMs) exhibit significant data inefficiency compared to human learners. A child is able to master language while consuming less than 100 million words of input, while language models require orders of magnitude more tokens during training. Our submission to the BabyLM Challenge utilizes a combination of self-distillation and reverse-distillation to train a sequence of ensemble models with improved training characteristics on a fixed-size 10 million-word dataset. Self-distillation is used to generate an ensemble of models of a certain fixed size, while reverse distillation is used to train a more expressive larger model from a previously trained generation of relatively smaller models, while largely preserving learned accuracy.We find that ensembles consisting of two smaller models and one identical born-again model serve as ideal ensembles for each trained generation of model size. We demonstrate that, although our method is not novel, it provides consistent and modest performance improvements on the BLiMP and GLUE benchmarks.

## related work knowledge distillation @xcite , a technique central to our work, has emerged as a popular approach for transferring knowledge from large models to smaller, more efficient ones. @xcite introduced the concept of "born again neural networks," where neural networks are trained using the predictions of an already-trained model, illustrating the potential of self-distillation. @xcite provided a comprehensive survey of various knowledge distillation techniques, categorizing them based on model types and applications and demonstrating their use in optimizing neural networks for various tasks, including language modeling. we build on work by @xcite , which contributed to the area by exploring knowledge distillation from an ensemble of teacher models trained on small datasets, achieving competitive results without performance degradation. whereas babyllama compressed large models into a smaller model, we attempt to use born-again ensembles of these smaller models to learn successively larger models. we find our techniques largely preserve and improve the base model's accuracy. while babyllama compressed model outperforms its teachers, our model expansion preserves these gains and allows us to continue learning with larger models. the larger expanded models have also been found to be more amenable to fine-tuning downstream tasks. we trained a series of decoder models with increasing sizes-58m, 95m, and 360m-following the training protocol outlined in figure 1 . each model size includes a sequence of three models, all based on the decoder-only llama architecture @xcite . the architectural details for each model variant are summarized in sequence zero for a given model size is trained using a teacher ensemble, which consists of three models strictly smaller than the current model. sequence one is trained with two smaller models and one model of the same size. sequence two is trained with two models of the same size and one smaller model. for each model size from 95m onward, three teacher models are used. however, the initial 58m model is trained in a strictly born-again sequence. our base model is 58 m sequence 0 is the base model @xcite , which we trained using this year's dataset from scratch. (note that this starting model performs below the later released contest baseline babyllama model). we apply identical prepossessing and tokenization as babyllama model on the 10 million word dataset, provided by babylm challenge organizers. ## results we evaluated the models on three benchmarks: glue @xcite , blimp @xcite , and ewok @xcite . for the glue benchmark, an additional fine-tuning phase was included to enhance the model's task-specific performance. detailed results are provided in appendix a. ## conclusion in this study, we have shown that we can train an ensemble of born-again teacher networks and use the ensemble of teachers to train larger student models. we find that having a model of the same size while having two models of smaller sizes in the ensemble leads to consistent improvements in the blimp benchmark. similar improvements are also noted on glue benchmarks, which included an intermediate finetuning step. we note that the accuracy of a smaller model is not lost in the reverse distillation process, thus allowing us to continue training with a larger models. for several of the benchmark tasks, however, we observe that improvements are non-monotonic but trend upward. thus, knowledge-distillation for student models is not consistently noise-free. this self-distillation and reverse-distillation process can be repeated to grow the size of our ensembles. with larger models more amenable to finetuning. further work is needed to quantify the limits of this method of improvement compared to directly training a large network and distilling it down to a smaller model. moreover, further work is required to quantify measures of catastrophic forgetting, as validation loss is often not predictive of benchmark performance and particular sub-task/skill.