Strong Teacher Not Needed? On Distillation in LLM Pretraining

[arxiv] Taiming Lu, Zhuang Liu
Princeton University

---

Knowledge distillation generally assumes a strong-to-weak relationship where stronger teachers yield better students. In this work we examine this assumption in large language model pretraining. By varying architecture sizes and training token budgets, we create strong-to-weak, same-level, and weak-to-strong teacher–student relationships, and study distillation's effectiveness under each with different mixes between language-modeling loss and distillation loss. Three findings emerge:

1. With proper loss mixing, weak-to-strong and same-level distillation improves over standard pretraining — even small and undertrained teachers benefit large students.
2. Making the teacher stronger can lead to saturated or even reversed gains.
3. Distillation improves generalization (out-of-distribution, downstream) more readily than in-domain fitting.

Our results challenge the common belief that distillation pretraining requires a strong teacher.

Prerequisites

• Hardware — a Google Cloud TPU VM pod. Paper runs use v6e-64 and v4-128; smaller topologies work after adjusting per_device_batch_size in the generator templates.
• Environment — clone this repository into ~/maxtext on the coordinator host and install dependencies into a virtualenv at ~/maxtext_env:

  git clone <REPO_URL> ~/maxtext
  cd ~/maxtext
  python3.10 -m venv ~/maxtext_env
  source ~/maxtext_env/bin/activate
  pip install -r requirements.txt

• Data — pretokenize a FineWeb-Edu split into array_record shards with quickstart/stream_fineweb_to_tokens.py.
• Runtime env vars — export before running any script:

  export BUCKET_NAME=<your-gcs-bucket>
  export TPU_PREFIX=<your-tpu-prefix>
  export WANDB_API_KEY=<your-wandb-key>   # optional; omit to disable wandb

Pipeline

1. Pretrain the teacher

Generate a per-configuration job script, then launch it:

python scripts/pretrain/generate_scripts.py \
    --archs 1b --tokens 50B --seeds 42 \
    --output-dir scripts/pretrain
bash scripts/pretrain/llama1b-finewebedu-vanilla-s42-50b.sh

Full checkpoints are written to gs://$BUCKET_NAME/ckpts/pretrain/<run_name>/checkpoints/<step>/items.

2. Convert the teacher to a param-only checkpoint

Strip optimizer state and downcast to bfloat16 so the teacher can be loaded as a frozen distribution during student training. The run_name below must match the naming convention expected by the distillation scripts: llama{arch}-vanilla-{tokens}-s{seed}.

python -m MaxText.generate_param_only_checkpoint MaxText/configs/base.yml \
    model_name=llama3.1-1b \
    load_full_state_path=gs://$BUCKET_NAME/ckpts/pretrain/llama3.1-1b-finewebedu-vanilla-s42-50b/checkpoints/24999/items \
    base_output_directory=$HOME/ckpts/pretrain_param_only \
    run_name=llama1b-vanilla-50B-s42

3. Distill into the student

Pick the experiment to reproduce. Each generator defaults to the paper's sweep; narrow with flags.

# Experiment 1: vary teacher size at fixed teacher tokens (50B), sweep alpha
python scripts/exp1/generate_scripts.py --output-dir scripts/exp1
bash scripts/exp1/exp1_llama1b-A1BT50BS42-a05-s43.sh

# Experiment 2: vary teacher training compute at fixed teacher size, sweep alpha
python scripts/exp2/generate_scripts.py --teacher-archs 1b --output-dir scripts/exp2
bash scripts/exp2/exp2_llama1b-A1BT80BS42-a05-s43.sh

Full sweep definitions and env-var documentation live in scripts/README.md.

Implementation

The distillation loss, teacher-distribution filtering (top-$k$ / top-$p$ with an optional "OTHER" bucket), temperature scaling, and hard-label variants are implemented on top of MaxText in MaxText/train.py and MaxText/max_utils.py.

• Pretrain — Vanilla CE pretraining for teacher checkpoints
• Experiment 1 — Distillation ablation varying teacher size
• Experiment 2 — Distillation ablation varying teacher training compute
• Quickstart — Single-job reproductions and FineWeb-Edu data preparation

Acknowledgement

This repository is built on top of MaxText, Google's high-performance JAX/TPU training framework for large language models. We gratefully acknowledge the generous support of the Google TPU Research Cloud (TRC), which provided the computational resources used to train every model in this study.

License

Apache 2.0, inherited from MaxText. See the LICENSE file for details.

Citation

@misc{lu2026strongKD,
  title={Strong Teacher Not Needed? On Distillation in LLM Pretraining},
  author={Lu, Taiming and Liu, Zhuang},
  year={2026},
}