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

TL;DR

This paper explores the extension of knowledge distillation to the pre-training phase of large language models, systematically analyzing various design choices to improve smaller models' performance.

Contribution

It introduces pre-training distillation (PD) for LLMs and systematically explores its design space, providing insights for better distillation practices.

Findings

Larger student LLMs benefit more from pre-training distillation.

A bigger teacher LLM does not always lead to better student performance.

Systematic exploration identifies optimal configurations for PD.

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.