Pre-training Distillation for Large Language Models: A Design Space Exploration

TL;DR

This paper systematically explore the design space of pre-training distillation across four aspects: logits processing, loss selection, scaling law, and offline or online logits, to find better configurations and interesting conclusions.

Abstract

Knowledge distillation (KD) aims to transfer knowledge from a large teacher model to a smaller student model. Previous work applying KD in the field of large language models (LLMs) typically focused on the post-training phase, where the student LLM learns directly from instructions and corresponding responses generated by the teacher model. In this paper, we extend KD to the pre-training phase of LLMs, named pre-training distillation (PD). We first conduct a preliminary experiment using GLM-4-9B as the teacher LLM to distill a 1.9B parameter student LLM, validating the effectiveness of PD. Considering the key impact factors of distillation, we systematically explore the design space of pre-training distillation across four aspects: logits processing, loss selection, scaling law, and offline or online logits. We conduct extensive experiments to explore the design space of pre-training distillation and find better configurations and interesting conclusions, such as larger student LLMs generally benefiting more from pre-training distillation, while a larger teacher LLM does not necessarily guarantee better results. We hope our exploration of the design space will inform future practices in pre-training distillation.

Figures (5)

• Figure 1: Results of the pre-trained 1.9B, 3.8B, and 6.8B student LLMs, using only LM loss, vanilla PD configuration (\ref{['sec:preliminary_exp']}), and a better PD configuration (PD$^*$) after our exploration. Details are placed in \ref{['sec:app_better']}.
• Figure 2: Relative improvements compared to LLM-LM using different $p$ in top-$p$-$100$ logits truncation and logits sizes per token with different $p$. The sizes are estimated using $10$ million tokens.
• Figure 3: Relative improvements compared to LLM-LM using different $k$ in top-$0.95$-$k$ logits truncation and logits sizes per token with different $k$.
• Figure 4: Relative improvements compared to LLM-LM using varying sizes of student and teacher LLMs.
• Figure 5: Experimental results of the checkpoints saved every $10,000$ step (about $83$B tokens) during the pre-training of 1.9B and 3.8B LLMs on 500B tokens. The last data point is from the checkpoint saved at the end.