Knowledge Distillation vs. Pretraining from Scratch under a Fixed (Computation) Budget (2024)

Abstract: Compared to standard language model (LM) pretraining (i.e., from scratch), Knowledge Distillation (KD) entails an additional forward pass through a teacher model that is typically substantially larger than the target student model. As such, KD in LM pretraining materially slows down throughput of pretraining instances vis-a-vis pretraining from scratch. Scaling laws of LM pretraining suggest that smaller models can close the gap to larger counterparts if trained on more data (i.e., processing more tokens)—and under a fixed computation budget, smaller models are able to process more data than larger models. We thus hypothesize that KD might, in fact, be suboptimal to pretraining from scratch for obtaining smaller LMs, when appropriately accounting for the compute budget. To test this, we compare pretraining from scratch against several KD strategies for masked language modeling (MLM) in a fair experimental setup, with respect to amount of computation as well as pretraining data. Downstream results on GLUE, however, do not confirm our hypothesis: while pretraining from scratch performs comparably to ordinary KD under a fixed computation budget, more sophisticated KD strategies, namely TinyBERT and MiniLM, outperform it by a notable margin. We further find that KD yields larger gains over pretraining from scratch when the data can be repeated under the fixed computation budget.

## introduction knowledge distillation (kd; @xcite during lm pretraining has emerged as the primary mean of compressing the capabilities of a large pretrained teacher model into a task agnostic smaller student model. kd is praised for yielding high-performing task agnostic small models, mitigating the need for pretraining contribution motivated by the recent findings in the realm of scaling laws and recognizing the absence of a fair comparison between kd and no-kd, our primary contribution lies in the comparison of no-kd against kd strategies for mlm while ensuring a fair setup with regards to compute budget and pretraining data. we initially assess no-kd in an optimal setup, where unlimited pretraining tokens are available within a fixed compute budget. additionally, we examine the scenario when data is constrained within a fixed compute budget. our results reveal that, in the optimal setting, no-kd performs indeed comparably to vanilla-kd, exhibiting an average improvement over vanilla-kd of 0.4 and 0.1 points for 6-, and 12-layer models on glue. however, no-kd falls short of surpassing more advanced kd strategies, exemplified by the comparison with tinybert and minilm. when available data is limited within the fixed compute budget, kd strategies outperform no-kd by an even larger margin: no-kd, though faster, needs more epochs, whereas kd strategies extract more information from limited data. ## distillation strategies vanilla-kd vanilla-kd for mlm pretraining is set up as follows. a small mlm student is trained to mimick the predictions for a particular training instance of a large pretrained mlm teacher: the distillation objective is to minimize the soft cross-entropy between the logits z t of the mlm teacher and the logits z s of the mlm student, with a temperature factor t: l @xmath0, z s /t). following @xcite , the final training loss equally combines l pred with the mlm loss l ce during pretraining. tinybert jiao et al. (2020) distill knowledge by minimizing the mean-squared error (mse) between latent representations of the mlm student s and the mlm teacher t by model layers as follows. first, the embedding matrices of the student (e s ) and the teacher (e t ) are aligned by minimizing the loss l @xmath1, e t ). the authors further fit the unnormalized attention scores per head h of the mlm student s to the mlm teacher t by optimizing lastly, the output hidden states h s of transformer layers of the student are also regressed onto the corresponding teacher output representations h t by optimizing 2020 ) also mimic the selfattention modules of the mlm teacher. unlike tinybert, minilm focuses on the last attention module. @xcite minimize the kl-divergence between the self-attention distributions of the mlm teacher and the mlm student. they further minimize the kl-divergence between the value relations of the mlm teacher t and mlm student s, i.e. ## discussion while our study provides insights into a fair evaluation of no-kd and kd for encoder-only models of moderate sizes, revealing negative results for no-kd, it may not cover the full spectrum of model sizes and architectures. for instance, @xcite show that for large decoder-only language models, no-kd performs comparably to vanilla-kd, aligning with our findings. however, advanced kd strategies like minilm exhibit poorer performance than no-kd and vanilla-kd, challenging both our results and common beliefs about kd regarding large decoder models. this disparity underscores the need for further investigation into a fair kd evaluation across a range of architectures and scales. additionally, we recommend investigating the impact of the teacher budget on performance in the fair setting, a consideration not closely examined in our current work. ## conclusion in this work, we investigate our hypothesis that, provided a fair training scenario, model pretraining from scratch yields similar results as kd during pretraining. our rationale is grounded in recent advancements in scaling laws for language models and that the literature lacks a fair comparison between no-kd and kd. our findings demonstrate that our initial assumption does not hold true: while, in an optimal setting for no-kd, no-kd performs on par with ordinary kd, it falls short when compared to more sophisticated kd strategies.