PaperlayerView exact revision on arXiv

Source-linked AI summary

MiniPLM: Knowledge Distillation for Pre-Training Language Models

Yuxian Gu, Hao Zhou, Fandong Meng, Jie Zhou, Minlie Huang

arXiv:2410.17215v3cs.CL

TL;DR

Pre-training KD faces a trade-off: online teacher inference adds costs and can constrain model-family flexibility, while offline approaches risk losing data difficulty and diversity. MiniPLM refines the training distribution offline using differences between large and small LMs. Experiments report gains across downstream tasks and language modeling, with reduced pre-training computation and data demand.

  • Problem

    Pre-training KD methods face efficiency and flexibility costs from online teacher inference and difficulty preserving diversity in offline-generated data.

    Source: §1 · p. 2

  • Method

    MINIPLM uses offline Difference Sampling based on teacher and small reference LM probabilities to refine the pre-training corpus for student training.

    Source: §1 · p. 2, §2.2 · p. 4, §2.5 · p. 5

  • Results

    MINIPLM improves zero-shot performance on 9 downstream tasks, enhances language modeling, and reduces pre-training computation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    MINIPLM reduces data demand by 2.4 times, mitigating the quick exhaustion of web-crawled corpora.

    Source: §1 · p. 2

    MINIPLM requires large-LM probabilities on pre-training-corpus texts, making black-box KD for close-source LMs challenging; obtaining probabilities through some APIs can require many calls.

    Source: §5 · p. 10

Abstract

from arXiv · show

Knowledge distillation (KD) is widely used to train small, high-performing student language models (LMs) using large teacher LMs. While effective in fine-tuning, KD during pre-training faces efficiency, flexibility, and effectiveness issues. Existing methods either incur high computational costs due to online teacher inference, require tokenization matching between teacher and student LMs, or risk losing the difficulty and diversity of the teacher-generated training data. In this work, we propose MiniPLM, a KD framework for pre-training LMs by refining the training data distribution with the teacher LM's knowledge. For efficiency, MiniPLM performs offline teacher inference, allowing KD for multiple student LMs without adding training costs. For flexibility, MiniPLM operates solely on the training corpus, enabling KD across model families. For effectiveness, MiniPLM leverages the differences between large and small LMs to enhance the training data difficulty and diversity, helping student LMs acquire versatile and sophisticated knowledge. Extensive experiments demonstrate that MiniPLM boosts the student LMs' performance on 9 common downstream tasks, improves language modeling capabilities, and reduces pre-training computation. The benefit of MiniPLM extends to larger training scales, evidenced by the scaling curve extrapolation. Further analysis reveals that MiniPLM supports KD across model families and enhances the pre-training data utilization. Our code, data, and models can be found at https://github.com/thu-coai/MiniPLM.

1 INTRODUCTION

MiniPLM addresses the difficulty of applying knowledge distillation during pre-training: online methods add teacher-inference costs, while offline methods can lose data difficulty and diversity. It uses offline Difference Sampling to refine the training distribution, and experiments report gains across downstream performance, language modeling, and pre-training data use.

1 INTRODUCTION

  • Small LMs offer deployment-friendly inference, but their training faces computation challenges because they are typically far from compute-optimal configurations.This has motivated interest in the limits of small LMs under pre-training-computation constraints.Source: §1 · pp. 1–2
  • At controlled training FLOPs, fine-tuning KD methods perform similarly to or worse than pre-training without KD, despite online KD gains at fixed steps.Online KD adds teacher-inference overhead and often requires shared tokenization, while offline KD can struggle to preserve difficulty and diversity.Source: Figure 2 · p. 2, §1 · p. 2
  • Difference Sampling uses offline teacher inference and differences between a large teacher LM and a small reference LM to refine the training distribution.It down-samples easy and common patterns, up-samples hard and diverse instances, and filters noisy or harmful data; students are then pre-trained from scratch on the refined distribution.Source: §1 · p. 2
  • MINIPLM improves zero-shot performance on 9 downstream tasks, enhances language modeling, and reduces pre-training computation.Its benefits remain consistent in scaling-curve extrapolation to ~10T tokens, extend across model families, and reduce data demand by 2.4 times.Source: §1 · p. 2

2 MINIPLM: KD FOR PRE-TRAINING LMS

MiniPLM uses teacher knowledge to refine the pre-training corpus offline through Difference Sampling, then trains students from scratch on that corpus. It frames distillation as rewarding teacher-preferred text while preserving diversity through comparison with a small reference LM.

2.1 KD AS REWARD MAXIMIZATION

  • MINIPLM formulates reverse-KL distillation as reward maximization, where teacher-preferred text earns high reward while low student probability supports diversity.The reward is based on the teacher-to-student probability ratio.Source: §2.1 · p. 3
  • Best-of-N samples M candidates from the student LM and retains the K highest-reward instances, but requires student sampling and reward computation and may not preserve candidate diversity or difficulty.The paper contrasts this with Difference Sampling, which avoids student-specific candidates and careful prompt engineering.Source: §2.1 · p. 3

2.2 DIFFERENCE SAMPLING

  • Difference Sampling selects high-reward instances directly from corpus D, and uses a small reference LM to approximate student difficulty and make the sampled corpus reusable.The method replaces student-based reward computation with the reference LM; the paper motivates this approximation by noting that relative data difficulty generally remains consistent across models.Source: §2.2 · p. 4

2.3 PRE-TRAINING ON DIFFERENCE-SAMPLED CORPUS

  • MINIPLM pre-trains students from scratch on the difference-sampled corpus using standard cross-entropy loss for next-token prediction.The loss averages token-level negative log probabilities over sequences in D′.Source: §2.3 · p. 4

2.4 MINIPLM TRAINING PIPELINE

  • The training pipeline samples a small reference subset, trains the reference LM, builds D′ from the remaining corpus using Difference Sampling, then trains students on D′.The sampling ratio α controls D′ size, with K = α|D − Dref|.Source: §2.4 · p. 4

2.5 DISCUSSION

  • MINIPLM stores teacher and reference probabilities offline, enabling reuse of the refined corpus across students without extra training computation or constraints on student architecture or tokenization.The paper reports 200MB of storage for 50B tokens at a 1,024 sequence length, compared with 30PB for storing per-token distributions for online KD.Source: §2.5 · p. 5
  • Difference Sampling down-samples common easy patterns, up-samples hard and diverse knowledge, and discards noisy instances that the teacher scores below the reference LM.Without comparing the teacher and reference LM, the corpus would lose the down-sampling of common patterns and could be dominated by them.Source: §2.5 · p. 5

3 EXPERIMENTS

Experiments compare MINIPLM with pre-training and KD baselines under controlled computation, then test its performance across tasks, model families, data budgets, and teacher sizes. The results report stronger downstream and language-modeling performance, reduced computation, and benefits in cross-family and data-limited settings.

3.1 EXPERIMENTAL SETUP

  • Experiments use a 1.8B Qwen-1.5 teacher and 200M, 500M, and 1.2B Qwen students, with a maximum of 50B training tokens.The setup compares MINIPLM with no-KD pre-training, Vanilla KD, SeqKD, and MiniLLM; online methods use fewer steps to match total computation.Source: §3.1 · pp. 5–6

3.2 MAIN RESULTS

  • MINIPLM achieves the best performance across most downstream tasks, and its improvements over Vanilla KD scale with student size.Table 1 reports zero-shot accuracy on nine tasks and average scores for 200M, 500M, and 1.2B students under matched training computation.Source: Table 1 · p. 6, §3.2 · pp. 6–7
  • MINIPLM matches Vanilla KD’s 500M-student performance while reducing computation by 2.2 times; at equal FLOPs, it also yields the lowest DCLM test losses.The language-modeling comparison uses the DCLM subset and Qwen-family students of 200M, 500M, and 1.2B parameters.Source: §3.2 · p. 7, Figure 4 · p. 7

3.3 KD ACROSS MODEL FAMILIES

  • Using Qwen-family teacher and reference models, MINIPLM outperforms no-KD pre-training and SeqKD for 212M Llama3.1 and 140M Mamba students.The table reports average downstream zero-shot accuracy and DCLM loss for these cross-family transfers.Source: §3.3 · pp. 7–8, Table 3 · p. 8

3.4 DATA-LIMITED SETTING

  • With a fixed 50B-token corpus, sampling 25% of the data and training for four epochs gives the best reported performance; the baseline is estimated to need 68B additional tokens to match it.The authors estimate this setting reduces the pre-training data requirement by 2.4 times; lower sampling ratios lead to rapid over-fitting.Source: §3.4 · p. 8

3.5 ANALYSIS

  • For a 200M student, a 500M teacher is most effective for MINIPLM, while neither smaller nor larger teachers are necessarily more helpful.The authors attribute weaker results with a 300M teacher to difficulty identifying valuable hard data, and with a 4B teacher to a scale mismatch that can make sampling rely on the reference LM.Source: §3.5 · pp. 8–9, Figure 6 · p. 9
  • Difference-sampled data has the highest measured semantic diversity among the compared corpora, which the authors associate with down-sampling repeated easy content and up-sampling hard, diverse texts.They argue that this offsets diversity lost when noisy parts of the corpus are removed.Source: §3.5 · p. 9, Table 4 · p. 9
  • Combining MINIPLM with Vanilla KD further improves performance for 500M and 1.2B students at matched FLOPs, with sequence-level and token-level teacher signals playing complementary roles.MINIPLM supplies coarse-grained sequence-level knowledge through pre-training data; Vanilla KD aligns token-level probabilities.Source: §3.5 · p. 9

4 RELATED WORK

Prior work improves pre-training through data curation, optimization, and model or objective changes, while small-model training remains computationally challenging. This work investigates KD as a complementary approach for building strong small base LMs.

4 RELATED WORK

  • Pre-training improvements include data curation, optimization strategies, and changes to model architectures or training objectives.The paper positions its use of existing LM knowledge alongside these approaches.Source: §4 · p. 10
  • Small LMs remain challenging to train because matching large-model capabilities can require computation beyond Chinchilla-optimal scaling [Hoffmann et al., 2022; Kaplan et al., 2020].Prior approaches include data-quality improvement and model pruning; the paper explores KD as complementary.Source: §4 · p. 10
  • Earlier generative-LM KD work includes training on large-model-generated text [Peng et al., 2023; Hsieh et al., 2023], while other methods target fine-tuning.Because pre-training establishes core capabilities, this work investigates KD during pre-training.Source: §4 · pp. 9–10

5 CONCLUSION

The conclusion presents MINIPLM as a way to refine pre-training data using differences between teacher and reference LMs, with reported gains across tasks and model families. It also identifies probability access and teacher-size scaling as constraints, and points to larger-than-teacher students as future work.

5 CONCLUSION

  • Using a 1.8B teacher for 200M, 500M, and 1.2B students, MINIPLM improves performance on 9 downstream tasks, language modeling capability, and pre-training computation.The conclusion also reports improved use of limited pre-training data and distillation across different model families.Source: §5 · p. 10
  • MINIPLM requires teacher probabilities on corpus texts, making black-box KD for closed-source LMs challenging; one API workaround entails many calls.The workaround uses user-provided softmax bias to obtain token probabilities one by one.Source: §5 · p. 10
  • Future work could apply the difference-sampled corpus to students larger than the teacher to explore weak-to-strong generalization.The paper suggests improved data difficulty and diversity may benefit such larger LMs.Source: §5 · p. 10

A PROOF OF PROPOSITION 2.1

The proof establishes that independently sampled sets increasingly agree on their top-K elements under a shared reward ranking. It first treats the maximum case, then extends the argument to top-K subsets.

A PROOF OF PROPOSITION 2.1

  • For K = 1, the proof reduces the claim to selecting the maximal-reward instance and shows the respective maximum probabilities converge to 1 as sample sizes grow.For other instances, the corresponding maximum probabilities converge to 0.Source: A PROOF OF PROPOSITION 2.1 · p. 18
  • For K > 1, independence and decomposition let the proof extend the matching argument from the maxima to each of the K highest-ranked elements.The proof uses equality of the top-K subsets to require corresponding ranked elements to match.Source: A PROOF OF PROPOSITION 2.1 · pp. 18–19
  • The probability that the two samples have the same top-K subset under r(x) approaches 1 as N and M grow.The proof concludes this result in Eq. (14).Source: A PROOF OF PROPOSITION 2.1 · p. 19

B MORE EXPERIMENTAL DETAILS

The appendix specifies model, training, and evaluation settings, including compute-aligned training controls for KD baselines. Its resource discussion reports that MINIPLM matches no-KD training time and memory, unlike the higher costs of online Vanilla KD.

B MORE EXPERIMENTAL DETAILS

  • The main experiments use batch size 512, sequence length 1,024, and 100K steps for roughly 50B tokens; MiniLLM and Vanilla KD steps are set to align training computation with no-KD pre-training.The models use AdamW with β1 = 0.9, β2 = 0.98, and 0.1 weight decay; learning rate warms up for 2K steps then decays by cosine scheduling.Source: B MORE EXPERIMENTAL DETAILS · p. 19, Table 7 · p. 20
  • LLaMA3.1 and Mamba cross-family experiments use the same pre-training data and optimization hyperparameters as the main Qwen experiments.Their model configurations are listed in Table 6.Source: B MORE EXPERIMENTAL DETAILS · p. 19
  • Zero-shot evaluation uses LM-Eval-Harness across the listed downstream datasets, while language modeling evaluation uses 10K DCLM documents.The downstream set includes Hellaswag, LAMBADA, Winograde, OpenbookQA, ARC-Easy/Challenge, PIQA, SIQA, and StoryCloze.Source: B MORE EXPERIMENTAL DETAILS · p. 19
  • The appendix lists model configurations and learning rates in Table 6, and training-step settings for Vanilla KD and MiniLLM in Table 7.Table 7 sets their training-time computation to match pre-training without KD, with a 1.8B teacher and 100K no-KD steps.Source: Table 6 · p. 20, Table 7 · p. 20

C DISCUSSION ON RESOURCE CONSUMPTION

  • MINIPLM requires the same training time and memory as pre-training without KD, whereas Vanilla KD uses significantly more runtime and GPU memory because of online teacher inference.Difference Sampling is performed once per teacher for any student models to be distilled; teacher and reference inference account for most of its runtime.Source: C DISCUSSION ON RESOURCE CONSUMPTION · p. 20

D MORE RESULTS

For 200M and 1.2B students, MINIPLM saves pre-training computation while outperforming Vanilla KD. Scaling-law extrapolation estimates that, in the data-limited setting, pre-training without KD needs 68B extra tokens to match MINIPLM trained on 50B tokens.

D.1 COMPUTATION SCALING CURVES OF MORE SIZES

  • MINIPLM saves pre-training computation and constantly outperforms Vanilla KD for 200M and 1.2B student LMs.The scaling curves compare average downstream-task accuracy against pre-training FLOPs.Source: §D.1 · p. 20

D.2 TEST LOSSES EXTRAPOLATION WITH SCALING LAWS

  • For the data-unlimited setting, the authors fit scaling-law curves on DCLM test losses and predict losses at 1T and 10T tokens.Vanilla KD is evaluated at the same training FLOPs as the other methods.Source: §D.2 · pp. 20–21
  • 68B extra training tokens are estimated to be needed for Pre-Train w/o KD to match MINIPLM trained on 50B tokens.The estimate comes from a Data-Constrained Scaling Law extrapolation; the corresponding token count is 118B for Pre-Train w/o KD.Source: §D.2 · p. 21

D.3 ABLATION STUDY

The ablations examine how reference-model size and difference-sampling ratio affect MINIPLM's downstream performance. Larger reference models improve performance with diminishing gains, while smaller sampling ratios yield better accuracy but require more source data to retain 50B sampled tokens.

Impact of the Reference Model

  • Figure 8 shows that larger reference models lead to better performance for language models trained with MINIPLM, but the improvement saturates as the reference LM grows larger.Source: §D.3 · p. 21

Difference Sampling Ratio

  • Smaller difference-sampling ratios yield better downstream zero-shot accuracy, but require a larger original corpus to retain 50B sampled tokens.The experiments mainly use α = 0.5 to control computation overhead.Source: §D.3 · p. 21

D.4 CORRELATIONS BETWEEN log p(x) pREF(x) AND log p(x)

The analysis tests whether a reference-model-based score can stand in for the student-model score when selecting training instances. The reported correlations and sampling accuracy support using the reference-model score in Difference Sampling.

  • The comparison uses reference models of 45M, 104M, and 200M parameters against scores from a 500M student LM.The reported measures are Pearson and Spearman correlation and sampling accuracy.Source: §D.4 · p. 22
  • Using log p(x) p_ref(x) achieves strong correlations with log p(x) q_θ(x) and high sampling accuracy, supporting its use in Difference Sampling.Sampling accuracy measures correct selected-or-discarded classifications at a 0.5 threshold, with selection using the student-model score as ground truth.Source: §D.4 · p. 22, Table 11 · p. 22

E DIFFERENCE SAMPLING WITH A PROXY MODEL

MINIPLM’s proxy variant estimates Difference Sampling values from a small corpus, reducing the need for teacher inference across the full corpus. The proxy variant uses much less inference computation than standard MINIPLM while retaining substantial improvement over Vanilla KD.

Motivation and Method

  • The proxy approach computes teacher and reference scores on a small subset, then fits a proxy model to predict scores across the corpus.A 0.1B-token proxy subset from a 50B-token corpus requires only 0.2% of the computation needed to compute scores on the full corpus.Source: E DIFFERENCE SAMPLING WITH A PROXY MODEL · pp. 22–23

Method

  • The method selects the K = α|D| instances with the highest inferred scores to form the difference-sampled corpus.The proxy model uses average reference-model hidden states and a linear head to predict the scores.Source: E DIFFERENCE SAMPLING WITH A PROXY MODEL · p. 23

Results

F CASE STUDY

The case study separates corpus instances by their relative fit to the teacher and reference LMs: valuable examples are up-sampled, while common patterns are down-sampled and noisy or context-poor data are discarded.

F CASE STUDY

  • Repeated HTML and LaTeX patterns are down-sampled rather than eliminated, while selected story dialogue remains to teach basic language skills.The study characterizes repeated patterns as easy for both LMs and the dialogue as relatively easy but still useful.Source: F CASE STUDY · pp. 23–24, Table 13 · p. 25
  • Long knowledge-rich documents, in-context translation examples, and commented code are retained as hard, valuable training instances.These examples highlight large-LM advantages that the student should learn from.Source: F CASE STUDY · p. 24, Table 14 · p. 26
  • Noisy or context-poor instances are discarded, although they account for only 6.3% of the corpus before Difference Sampling.Examples include irregular symbols, meaningless number strings, and contextless scientific-paper data; the reference LM can overfit surface patterns that the teacher recognizes as unnatural or unhelpful.Source: F CASE STUDY · p. 24, Table 15 · p. 27

Related papers

Loading 2410.17215v3…