Teaching Tiny Minds: Exploring Methods to Enhance Knowledge Distillation for Small Language Models (2024)

Abstract: In this paper, we build off of the success of the previous BabyLM challenge winner’s model, BabyLlama, to explore various methods of enhancing knowledge distillation for small language models. Our main focus is on investigating how small a language model can be while still maintaining competitive performance. We experiment with three main approaches: (1) DistilledGPT-44M, which uses smaller teacher models and a more compact student model compared to BabyLlama; (2) ContrastiveLlama-58M, which incorporates contrastive loss into the knowledge distillation process; and (3) MaskedAdversarialLlama-58M, incorporates adversarial loss into the knowledge distillation process. Using the 10M-word dataset from the BabyLM challenge’s strict-small track, we evaluate our models on the BLiMP, EWoK, and GLUE benchmarks. Our results show that effective knowledge distillation can still be achieved with significantly smaller teacher and student models. In particular, our model DistilledGPT-44M is able to achieve better performance than one of last year’s winning entries, LTG-BERT, while achieving similar performance but cutting training time by around 70% and parameters by around 25% compared to the other winning entry, BabyLlama.

## introduction since 2017, transformers have been everywhere in nlp @xcite . their non-autoregressive nature allows for high parallelization, leading to unprecedented scalability. in recent years, a number of models with trillions of parameters have emerged, such as google's switch transformer (1.6 trillion) and huawei's pangu-σ (1.1 trillion). models like these or models in the billions demand enormous computational resources and vast swathes of training data. they consume substantial energy, raising concerns about their environmental impact @xcite . the costs of contin-ued scaling are also increasingly prohibitive, highlighting the need for more sample-efficient model architectures. the 2024 babylm challenge @xcite , by limiting the amount of training data available, in some ways aims to address the computational concerns around large models. however, the scope of the contest focuses more on limiting the amount of training data rather than limiting parameter size or compute; participants have the freedom to use models as large as they want. but last year, the winners of babylm @xcite demonstrated with their model babyllama that a small model can outperform models close to an order of magnitude larger on nlp tasks. model parameter efficiency does not necessarily mean worse results; in fact, in some cases it means better results. babyllama used knowledge distillation and ensemble learning to distill knowledge from two teacher models -gpt2-705m and llama-360m -to a smaller llama-58m student model @xcite . as a model compression technique, knowledge distillation (kd) has several advantages: it only requires access to the teacher model's output logits (not its weights), and it is also model agnostic. building on babyllama's success, we aim to demonstrate that even smaller teachers and students can achieve competitive performance, further pushing the boundaries of parameter efficiency. we explore the impact of using teachers with fewer parameters and distilling knowledge into even smaller student models. we also explore incorporating different losses into the distillation training, such as contrastive loss and adversarial loss. our results suggest that effective knowledge distillation can be achieved with significantly smaller teacher and student models, demonstrating competitive performance even with reduced parameters. we find that our distilledgpt-44m model, despite having much fewer parameters, achieves results comparable to the original babyllama-58m on key benchmarks. our experiments with contrastive and adversarial learning techniques in the distillation process, which albeit less promising, reveal interesting trade-offs between different aspects of model performance. ## dataset we use the provided 10m dataset from the strictsmall track and build on the babyllama repository https://github.com/timinar/babyllama . following their preprocessing steps, we apply regexbased cleaning and train a byte-pair encoding tokenizer on the training set. the train and dev sets are split into 128-token chunks, with the model being presented a new random permutation of these chunks in each epoch. validation loss is computed at the end of each epoch using a fixed, randomly sampled subset of the dev set. ## evaluation evaluation of model performance was done using the babylm evaluation suite @xcite . this consists of the following benchmarks: ## results our results for these 3 experiments can be found in ## limitations and future work although we showed the effectiveness of knowledge distillation with smaller models, we did not thoroughly explore the lower bounds of model size. in future experiments we could investigate even smaller student models or experiment with a wider range of teacher-student size combinations to find the optimal balance between model size and performance. additionally, our experiments with contrastive and adversarial learning techniques (contrastivellama-58m and maskedadversarialllama-58m) did not show improvements over the simpler distilledgpt-44m model. these advanced techniques probably require further refinement or different implementation strategies to be effective: we could try different weighting schemes for loss components in contrastive and adversarial training. additionally, for the masked adversarial model, the performance of the generator plays a critical role in generating effective perturbed inputs. using a more powerful mlm generator, rather than the smaller electra-56m model we used, could improve the adversarial training process and create better perturbations. ## conclusion herein, we showed that knowledge distillation can be used even with two very simple parents with around the same number of parameters as the child model, to produce a child model which outperforms both parents. we present distillgpt-44m, which outperforms both the baseline (gpt2) and one of last year's winning entry for the babylm challenge ltg-bert, while maintaining comparable performance to the other winning entry babyllama-58m despite reducing number of parameters by around 25% and cutting training time by around 70%.