Research map / Label-efficient and robust learning / Knowledge distillation
MiniPLM: Knowledge Distillation for Pre-training Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MiniPLM, an offline knowledge distillation framework for pre-training that refines the pre-training corpus via Difference Sampling using the log-probability gap between a large teacher LM and a small reference LM.
Problem
Knowledge distillation during pre-training is costly with online teacher inference, requires tokenizer matching, and offline KD data generation loses data difficulty and diversity.
Approach
Reformulate KD as reward maximization and approximate it by top-K selection from the pre-training corpus using reward log p_teacher(x) - log p_ref(x), where the reference LM is a tiny model trained on a small uniform subset; the refined corpus is then reused to pre-train any student from scratch.
Key results
- With a 1.8B teacher, MiniPLM raises the 200M student's 9-task average zero-shot accuracy to 41.3 vs 39.9 for Pre-Train w/o KD and Vanilla KD.
- MiniPLM matches Vanilla KD's performance on the 500M student while reducing pre-training compute by 2.2x.
- In the data-limited setting, MiniPLM reduces the data requirement by 2.4x (estimated 68B additional tokens for the baseline to match).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MiniPLM Difference Sampling instead of Vanilla KD | Qwen-1.5 200M student · 9 downstream tasks (Avg.) | accuracy: 39.9 -> 41.3 | Helped |
| Use MiniPLM Difference Sampling instead of Vanilla KD | Qwen-1.5 500M student · 9 downstream tasks (Avg.) | accuracy: 43.6 -> 44.8 | Helped |
| Use MiniPLM Difference Sampling instead of Vanilla KD | Qwen-1.5 1.2B student · 9 downstream tasks (Avg.) | accuracy: 45.4 -> 46.6 | Helped |
| Use MiniPLM Difference Sampling instead of Pre-Train w/o KD | Qwen-1.5 students · DCLM, 200M student, L_1T | loss: 3.35 -> 3.28 | Helped |
| Use MiniPLM Difference Sampling instead of Vanilla KD | Qwen-1.5 students · DCLM, 500M student, L_1T | loss: 3.12 -> 3.06 | Helped |
| Use MiniPLM Difference Sampling instead of Vanilla KD | Qwen-1.5 students · DCLM, 1.2B student, L_1T | loss: 2.95 -> 2.92 | Helped |
| Add Vanilla KD to MiniPLM's difference-sampled corpus (ablation) | Qwen-1.5 200M student · 9 downstream tasks (Avg.) | accuracy: 41.3 -> 40.7 | Hurt |
| Add Vanilla KD to MiniPLM's difference-sampled corpus (ablation) | Qwen-1.5 500M student · 9 downstream tasks (Avg.) | accuracy: 44.8 -> 44.9 | Helped |
| Add Vanilla KD to MiniPLM's difference-sampled corpus (ablation) | Qwen-1.5 1.2B student · 9 downstream tasks (Avg.) | accuracy: 46.6 -> 48.1 | Helped |
| Set sampling ratio to α=0.25 and train for 4 epochs instead of α=0.5 or α=0.125 (ablation) | MiniPLM, 200M student · DCLM, data-limited setting | loss | Helped |
| Use a 500M teacher rather than smaller or oversized teachers (ablation) | MiniPLM, 200M student · Downstream tasks | average zero-shot accuracy | Helped |
| Use MiniPLM instead of Vanilla KD | 500M student · Downstream tasks at matched performance | training computation: 2.2 times | Helped |
Surprising findings
- Vanilla KD does not beat Pre-Train w/o KD for the 200M student under equal compute, and offline SeqKD performs similar or worse than no KD.
- A larger teacher (e.g., 4B) does not necessarily help MiniPLM, since the log-probability scale of the teacher becomes too small relative to the reference LM.
Limitations
- Requires access to teacher token probabilities over pre-training texts, so it cannot be applied directly to black-box closed-source LMs.
- Performance does not monotonically improve with teacher size; oversized teachers shrink the log-probability scale and degrade difference sampling.
- Evaluated only with a 1.8B teacher and models up to 1.2B parameters; benefits at larger scales are extrapolated.
Experimental setup
Datasets: The Pile DCLM subset
Benchmarks: HellaSwag LAMBADA WinoGrande OpenBookQA ARC-easy ARC-challenge PIQA SocialIQA StoryCloze DCLM subset language modeling loss
Metrics: zero-shot accuracy average accuracy test loss training FLOPs efficiency
Baselines: Pre-Train w/o KD Vanilla KD MiniLLM SeqKD
Builds on: Knowledge distillation (Hinton et al.) MiniLLM SeqKD Best-of-N sampling Qwen1.5
Closest papers by meaning
- DistiLLM: Towards Streamlined Distillation for Large Language Models (ICML 2024 · similarity 0.93)
- DDK: Distilling Domain Knowledge for Efficient Large Language Models (NeurIPS 2024 · similarity 0.91)
- Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models (NeurIPS 2026 · similarity 0.90)
- Distilled Pretraining: A modern lens of Data, In-Context Learning and Test-Time Scaling (ICLR 2026 · similarity 0.90)
- DA-KD: Difficulty-Aware Knowledge Distillation for Efficient Large Language Models (ICML 2025 · similarity 0.90)
- Why Knowledge Distillation Works in Generative Models: A Minimal Working Explanation (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.