PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hao Peng, Xin Lv, Yushi Bai, Zijun Yao, Jiajie Zhang, Lei Hou, Juanzi Li

arXiv:2410.16215v1cs.CLcs.AI

TL;DR

LLM knowledge distillation has mainly been studied after pre-training, leaving its feasibility and best configurations during pre-training underexplored. This paper introduces pre-training distillation and systematically evaluates its design space, finding that PD improves student performance while its benefits depend on model scale and configuration.

  • Problem

    LLM knowledge distillation has typically focused on post-training, leaving whether and how distillation should be applied during pre-training insufficiently explored.

    Source: §1 · pp. 1–2

  • Method

    The paper distills teacher-generated token logits into students during pre-training and studies logits processing, loss selection, scaling, and offline versus online logits.

    Source: §1 · pp. 1–2

  • Results

    PD improves student performance; larger students generally benefit more, while larger teachers do not necessarily yield better results.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The explored design space identifies configurations and empirical conclusions intended to inform future pre-training distillation practices.

    Source: §5 · pp. 8–9

    The study does not explore interactions among different PD factors because the required combinations are too resource-intensive.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The paper extends knowledge distillation from LLM post-training to pre-training, introducing pre-training distillation (PD) and systematically exploring its design space. Experiments show PD improves student performance, with effects depending on student size, teacher size, loss configuration, and logit availability.

Motivation and approach

  • PD transfers teacher-generated token logits to a student during pre-training, using distillation loss to provide richer supervision than one-hot labels.The teacher logits can serve as a form of label smoothing and may accelerate student training or improve performance.Source: §1 · p. 1

Design-space exploration

  • The study explores PD across logits processing, loss selection, scaling law, and offline versus online logits.These dimensions cover storage-efficient teacher-logit processing, distillation objectives, model and corpus sizes, and when logits are generated.Source: §1 · p. 2

Preliminary experiment

  • A preliminary experiment distilled a 1.9B student from GLM-4-9B logits for 100B tokens, improving average performance by 1.6% across English and Chinese datasets.Top-p-k truncation reduced logit storage by 4,000× to approximately 15 TB.Source: §1 · p. 2

Loss selection

  • Kullback–Leibler divergence and negative log-likelihood yield similar improvements, while MSE causes a significant performance drop.Combining LM and KD losses with WSD scheduling performs best among the tested loss configurations.Source: §1 · p. 2

Scaling law

  • Larger student LLMs generally benefit more from PD, whereas a larger teacher LLM does not necessarily produce better results.The authors suggest the latter pattern may reflect a capacity gap between teacher and student models.Source: §1 · p. 2

Offline or online

  • Online logits also improve PD performance, although less than offline logits, enabling logits to be saved during teacher pre-training without additional inference cost.The authors report that this can support PD for a series of smaller LLMs.Source: §1 · p. 2

2 Design Space for PD

The paper formulates pre-training distillation around teacher and student token distributions, then organizes its empirical study around four design dimensions. These dimensions vary teacher-logit processing, loss construction, model and corpus scale, and offline versus online teacher outputs.

Objective

  • The PD objective compares student token probabilities with transformed teacher probabilities over each token in a pre-training sequence.The distillation loss may use Kullback–Leibler divergence.Source: §2 · p. 2

Loss selection

  • Loss selection varies the distillation loss and the combination factor between traditional language-modeling loss and distillation loss.The study therefore treats both the KD objective and its mixture with LM loss as design choices.Source: §2 · pp. 2–3

Logits processing

  • Teacher logits are processed through truncation and temperature-based normalization before entering the distillation objective.The transformation F reduces or normalizes the teacher distribution, while τ controls normalization temperature.Source: §2 · p. 2

Scaling law

  • Scaling-law experiments vary student and teacher model sizes as well as the corpus size used for student pre-training.These experiments test how PD changes across model-capacity and data-scale regimes.Source: §2 · pp. 2–3

Offline or online

  • The teacher distribution can be obtained offline from a pre-trained teacher or online while the teacher is being pre-trained.The paper compares these strategies as distinct ways to supply PθT(x_t|x_<t).Source: §2 · pp. 2–3

3 Experiments

The experiments validate pre-training distillation and explore logits processing, loss selection, scaling behavior, and offline versus online teacher logits. Across these dimensions, PD generally improves student LLMs, with benefits shaped by truncation choices, loss mixtures, model-size gaps, and teacher-training stage.

3.1 Preliminary Experiment

  • LLM-KD generally outperforms LLM-LM after distilling a 1.9B student from GLM-4-9B, validating PD while indicating that the initial configuration is not optimal.The setup uses 100B tokens, top-0.95-100 truncation, and approximately 15TB storage after a 4,000× reduction.Source: §3.1 · p. 3

3.2 Design Dimension #1: Logits Processing

  • Smaller p values yield similar gains, while top-0.95-k truncation improves performance for all k and performs best at k = 50, enabling further storage savings.For k = 1, equivalent to teacher-generated labels with LM loss, improvement may reflect implicit noise filtering in the teacher.Source: §3.2 · p. 4, Figure 2 · p. 4, Figure 3 · p. 4
  • Lower temperatures τ ≤2.0 produce similar improvements, whereas τ ≥5.0 limits gains; AdaKDH performs best but does not significantly exceed static τ = 0.5.The results suggest that overly uniform teacher distributions are inefficient for student learning.Source: §3.2 · pp. 4–5, Table 2 · p. 4, Table 3 · p. 5

3.3 Design Dimension #2: Loss Selection

  • NLL and KLD outperform LLM-LM, KLD generally exceeds NLL, NLL is stronger on MMLU and C-Eval, and MSE causes a significant performance decline.This differs from image-classification distillation findings that favor MSE, indicating distinct LLM pre-training dynamics.Source: §3.3 · p. 5
  • The best static loss mixture uses α = 0.9, while WSD-α combined with WSD-LR performs best among dynamic schedules.The findings indicate that a small LM-loss share and substantial KD loss, especially during high-learning-rate training, can improve PD.Source: §3.3 · pp. 5–6, Table 4 · p. 5, Table 5 · p. 6

3.4 Design Dimension #3: Scaling Law

  • Larger students generally benefit more from PD, but larger teachers do not necessarily improve results because the teacher–student capacity gap can limit distillation.PD is effective when the student reaches about 10% or more of the teacher size, with benefits increasing until an unexplored turning point.Source: §3.4 · pp. 6–7
  • Compared with LM-only training, PD improves students throughout 500B-token pre-training, with gains increasing initially before converging with a slight decrease.The gains remain significant at the end of training, although trillion-level token regimes were not tested.Source: §3.4 · p. 7, Figure 5 · p. 7

3.5 Design Dimension #4: Offline or Online

  • Online logits from a non-converged teacher perform poorly, but later-stage logits with α = 0.1 and top-0.95-50 truncation slightly outperform LLM-LM while remaining below offline-logit PD.Offline logits are preferred for training one LLM, whereas online logits can support sequential training of students after a larger teacher.Source: §3.5 · pp. 7–8

4 Related Work

Prior work applies knowledge distillation broadly and commonly uses post-training teacher-generated responses for LLMs, while pre-training distillation research has been less systematic for billion-parameter models. This paper addresses that gap by systematically exploring pre-training distillation for LLMs.

  • Knowledge distillation transfers knowledge from large teacher models to smaller students and has been applied in computer vision, natural language processing, and speech recognition.Source: §4 · p. 8
  • Post-training KD commonly trains LLM students with teacher-forcing language-modeling loss on instructions and responses generated by advanced teacher LLMs such as GPT-4.Source: §4 · p. 8
  • Earlier pre-training distillation work targeted models with only several million parameters, so its configurations may not directly apply to billion-parameter LLMs.Source: §4 · p. 8
  • Recent LLM distillation studies include pruning and pre-training distillation, but several provide limited process details, while systematic exploration of pre-training distillation remains limited.Source: §4 · p. 8
  • This work systematically studies pre-training distillation’s key impact factors and seeks configurations applicable to prior pruning and distillation efforts.Source: §4 · p. 8

5 Conclusion

The paper systematically explores four dimensions of pre-training distillation and identifies configurations and scaling patterns that can inform future practice. Its conclusions emphasize stronger benefits for larger students, without a guaranteed advantage from larger teachers.

  • The study examines logits processing, loss selection, scaling law, and offline-versus-online strategies for obtaining logits in pre-training distillation.Source: §5 · pp. 8–9
  • Larger student LLMs generally benefit more from pre-training distillation, whereas larger teacher LLMs do not necessarily guarantee better results.Source: §5 · pp. 8–9
  • The exploration aims to identify better pre-training-distillation configurations and inform future practices.Source: §5 · pp. 8–9

Limitations

The study does not explore interactions among different pre-training-distillation factors because exhaustive combinations are too resource-intensive. The authors also note substantial computational costs and associated environmental impact.

  • The main limitation is that the study does not evaluate combinations or interactions among different pre-training-distillation factors.Source: Limitations · p. 9
  • Exhaustively searching factor combinations is considered unaffordable because the experiments are highly resource-intensive.Source: Limitations · p. 9
  • The controlled experiments already incurred significant computational costs that emitted substantial carbon dioxide and negatively impacted the environment.Source: Limitations · p. 9

Ethical Considerations

The paper describes measures concerning intellectual property, intended use, risk control, and AI assistance. It reports adherence to model and dataset licenses and states that the data used was properly anonymized.

A Experimental Details and more Results

This section reports experimental details and additional results, with all experiments conducted on Nvidia H800 GPUs.

A.1 Preliminary Experiment

The preliminary experiment specifies the student’s supervised fine-tuning and evaluation setup, including mixed training data, response-only loss, and task-specific few-shot settings.

  • The SFT data combine 10B instruction-tuning tokens with an additional 10B pre-training text corpus.Source: §A.1 · p. 13
  • SFT computes language-modeling loss only on the response portion of instruction-tuning examples.Source: §A.1 · p. 13
  • Evaluation uses task-specific settings: zero-shot for HellaSwag, WinoGrande, PIQA, and KBQA, and 5- to 8-shot evaluation for the remaining listed benchmarks.The settings are 5-shot for C3 and C-Eval, 6-shot for MMLU, and 8-shot for GSM8k.Source: §A.1 · p. 13

A.2 Logits Processing

The logits-processing implementation evaluates adaptive temperature methods using truncated logits and specified normalization parameters, with results reported across evaluation datasets.

  • Adaptive temperature calculation uses NormKD (Chi et al., 2023) and WTTM (Zheng and YANG, 2024).Source: §A.2 · p. 13
  • The implementation uses truncated logits rather than logits from the entire vocabulary.NormKD sets T_norm to 1.0 and α to 0.5, while WTTM sets γ to 0.1 and β to 1.0.Source: §A.2 · p. 13
  • For entropy-based temperature calculation, H denotes each sample’s entropy, Hmax is estimated on 10 million tokens, and τmax and τmin are set to 2.0 and 0.1.The implementation sets Hmax to 4.8.Source: §A.2 · p. 13

A.3 Loss Selection

The experiments examine loss scheduling, student and teacher model sizes, offline logits, and a selected configuration for pre-training distillation, with results reported across evaluation datasets.

A.3 Loss Selection

  • The WSD scheduler uses a linear schedule during warmup and a cosine schedule during decay, with experiments varying α across evaluation datasets.Source: §A.3 · p. 13, Table 10 · p. 14

A.4 Model Size

  • For model-size experiments, 1.9B and 3.8B students are pretrained on 500B tokens, checkpointed every 10,000 optimization steps, then SFT-trained before evaluation.A final checkpoint is also saved, and averaged performance is reported in Figure 5.Source: §A.4 · p. 13, Table 11 · p. 15, Table 12 · p. 15

A.5 Offline or Online

  • The 9B teacher is pretrained from scratch on 400B tokens while storing logits, requiring about 180TB of disk space.Because the teacher has not converged, its logits are described as more uniform and noisier.Source: §A.5 · p. 13

A.6 A Better Configuration for PD

  • The selected configuration uses top-0.95-50 truncation, τ = 2.0 normalization, KLD combined with LM loss, WSD-α and WSD-LR, GLM-4-9B, and offline logits.It distills 1.9B and 3.8B student LLMs, with the WSD maximum α set to 0.9.Source: §A.6 · p. 13, Table 13 · p. 16
  • The experiments compare truncation settings, normalization temperatures, teacher and student sizes, checkpoints, and the selected distillation configuration across evaluation datasets.These comparisons are reported in Tables 8–14, including the checkpoint, model-size, and final-configuration evaluations.Source: §A.6 · p. 13, Table 10 · p. 14, Table 8 · p. 14

Related papers

Loading 2410.16215v1…