Large Language ModelsInference & Performance Optimization

How does LLM distillation work?

LLM distillation works by training a student model to learn from the outputs or probability distributions of a stronger teacher model. The teacher provides supervision; the student updates its own parameters to reproduce useful teacher behavior.

Step-by-step workflow

  1. Choose a teacher and student. Select a capable teacher and a student small enough to meet the intended memory, latency, or deployment budget.
  2. Prepare representative prompts. Include examples resembling the real tasks the student must handle. Keep separate validation and test sets.
  3. Collect teacher supervision. The teacher may generate answers, labels, explanations, or next-token probability distributions. Which signals are accessible depends on the teacher and API.
  4. Train the student. Compare student predictions with teacher targets, calculate a loss, and update the student's weights using backpropagation. The teacher normally remains frozen.
  5. Evaluate and refine. Measure task accuracy, generalization, safety, latency, and cost on held-out examples. Improve data or training if necessary.
  6. Deploy the student. If its quality is sufficient, serve the student without running the teacher for every user request.
Task prompts ──→ Teacher model ──→ Teacher signals
      │                                   │
      └────────→ Student model ────────────┤
                                          ↓
                                    Training loss
                                          ↓
                                  Update student only

What loss does the student minimize?

There are two widely used forms of supervision:

  • Response-level distillation: Train on the teacher's generated answers, often with standard cross-entropy loss over the answer tokens. This works even when the teacher exposes only text.
  • Logit-level distillation: Use the teacher's next-token probabilities as soft targets, often with a KL-divergence-based objective. This preserves information about alternative tokens that a single generated answer omits.

Some methods combine these losses. When teacher and student use different tokenizers, directly aligning their token probabilities may require special handling.

Simple example

Suppose a teacher classifies a support message:

Teacher probabilities:
Billing:   0.70
Account:   0.25
Technical: 0.05

Student before training:
Billing:   0.50
Account:   0.35
Technical: 0.15

A distribution-matching loss encourages the student to move toward the teacher's distribution. The teacher communicates not only that Billing is the most likely answer, but also that Account is a plausible alternative.

For generative LLMs, this idea applies to next-token distributions, not just three classification labels.

Offline vs. on-policy distillation

  • Offline (off-policy): Train the student from a fixed dataset of teacher-generated responses or teacher scores. It is comparatively simple, but the student may encounter different continuations when generating on its own.
  • On-policy: Let the student generate completions, then have the teacher evaluate those continuations. This can help the student learn from its own mistakes and reduce training–inference mismatch.

For example, Hugging Face TRL's DistillationTrainer implements on-policy training that matches student and teacher token distributions.

What are the main limitations?

Distillation requires teacher inference and student training, may transfer teacher errors, and does not guarantee the student matches every teacher capability. Good results on a narrow task do not imply equivalent general intelligence.

See “What is model distillation for LLMs?” for the conceptual overview.

In short: collect teacher supervision, optimize the student's predictions against it, and test the smaller model on real workloads before deploying it independently.

References: Hugging Face DistillationTrainer, Distilling Step-by-Step!, and On-Policy Distillation of Language Models.

Sources & Research References

#how does LLM distillation work#LLM distillation step by step#teacher student LLM training#logit-level distillation KL divergence#on-policy LLM distillation#response level knowledge distillation#large-language-models#llm#knowledge-distillation#teacher-student

Learning Sequence

Large Language Models Curriculum

69% Completed
Up NextLarge Language Models

What is speculative decoding, and how does it accelerate LLM inference?

Explore Related Questions

Deepen your understanding across related Large Language Models topics.

Machinoai Logo
Machinoai

Your edge in the AI era. Deep dives into machine learning concepts, paper explainers, verified industry jobs, and developer tools.

© 2026 Machinoai. All rights reserved.