Distillation for inference

In the previous post, I laid out the tradeoff between small fine-tuned models and large quantized ones. Distillation is the bridge between them. The idea is simple: run a large, expensive model (the teacher) on your task distribution, collect its outputs, and train a small, cheap model (the student) to mimic them. You pay the teacher's compute cost once during training, then serve the student forever at a fraction of the cost.

This is not a new idea. Hinton, Vinyals, and Dean published the seminal "Distilling the Knowledge in a Neural Network" paper back in 2015. But the technique has become especially important for LLM inference because the cost gap between large and small models is so dramatic: a 70B model costs roughly 8x more per token to serve than an 8B model on the same hardware class.

Classical distillation: soft labels

The original distillation approach trains the student on the teacher's output probability distribution, not just the hard labels. For a classification task, the teacher might assign 70% probability to class A, 25% to class B, and 5% to class C. Training on these "soft labels" gives the student more signal than just "the answer is A," because it also learns which alternatives the teacher considered plausible.

The loss function combines two terms:

# Classical distillation loss
def distillation_loss(student_logits, teacher_logits, labels,
                      temperature=2.0, alpha=0.5):
    # Soft loss: KL divergence on temperature-scaled distributions
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
    soft_loss = soft_loss * (temperature ** 2)  # scale gradient

    # Hard loss: standard cross-entropy on ground truth
    hard_loss = F.cross_entropy(student_logits, labels)

    return alpha * soft_loss + (1 - alpha) * hard_loss

The temperature parameter controls how much the soft distribution is "smoothed." Higher temperatures make the teacher's distribution more uniform, which transfers more information about relative class similarities. A temperature of 2 to 4 is typical for LLM distillation.

LLM distillation in practice: output-based

For large language models, classical distillation is often impractical because it requires access to the teacher's full logit distribution at every token position. If your teacher is a 70B model, storing and transmitting those logits (vocabulary size is typically 32K to 128K) for every token in every training example is expensive.

The more common approach in practice is output-based distillation: you run the teacher on your inputs, collect the generated text, and fine-tune the student on (input, teacher_output) pairs using standard supervised fine-tuning. No logits needed.

# Output-based distillation pipeline
# Step 1: Generate teacher outputs
teacher = load_model("llama-3.1-70b")
student_training_data = []

for prompt in task_prompts:
    response = teacher.generate(prompt, temperature=0.7)
    student_training_data.append({
        "input": prompt,
        "output": response
    })

# Step 2: Fine-tune student on teacher outputs
student = load_model("llama-3.1-8b")
train(student, student_training_data, epochs=3, lr=2e-5)

This is sometimes called "synthetic data fine-tuning" rather than distillation, but the principle is the same: the teacher's knowledge is transferred to the student through the training data.

What knowledge transfers and what does not

Distillation reliably transfers:

• Output format and style: If the teacher produces well-structured JSON, step-by-step reasoning, or a specific writing style, the student learns to replicate it.
• Task-specific patterns: Classification decisions, extraction patterns, summarization approaches that the teacher applies consistently.
• Implicit preferences: How the teacher handles ambiguity, what level of detail it provides, which edge cases it handles gracefully.

Distillation struggles to transfer:

• Deep factual knowledge: If the teacher knows an obscure fact because it was in the training data, the student may not have the capacity or the pretraining coverage to retain it from a few examples.
• Multi-step reasoning: Complex chains of logic that depend on the teacher's larger hidden state. The student may learn the surface pattern of reasoning without the underlying capability.
• Robustness to distribution shift: The student learns to mimic the teacher on the training distribution. Novel inputs outside that distribution may produce poor results.

The distillation pipeline

Here is the pipeline I recommend for production distillation:

1. Define the task distribution: Collect or generate a representative set of inputs. This is the most important step. The student can only learn from examples it sees.
2. Generate teacher outputs: Run the teacher on all inputs. Use a moderate temperature (0.3 to 0.7) to get high-quality but slightly diverse outputs. Generate multiple outputs per input if you want to do rejection sampling.
3. Filter for quality: Not all teacher outputs are good. Use automated checks (format validation, consistency checks) or a separate evaluator model to filter out low-quality outputs.
4. Fine-tune the student: Standard supervised fine-tuning with LoRA or full fine-tuning. 3 to 5 epochs on the filtered dataset is usually enough.
5. Evaluate on held-out data: Compare the student against the teacher on a test set using your eval harness. Track both automated metrics and manual review.

Advanced: on-policy distillation

The pipeline above is "off-policy": you generate teacher data once and train on it. A more effective (but more expensive) approach is on-policy distillation, where the student generates outputs, the teacher scores or corrects them, and the student trains on the corrections.

# On-policy distillation loop
for epoch in range(num_epochs):
    for prompt in task_prompts:
        # Student generates
        student_output = student.generate(prompt)

        # Teacher corrects or scores
        teacher_output = teacher.generate(
            f"Improve this response:\n{student_output}\n"
            f"Original prompt: {prompt}"
        )

        # Train student on teacher's correction
        train_step(student, prompt, teacher_output)

This is related to RLHF and DPO, but uses the teacher model as the reward signal instead of human preferences. It tends to produce better results because the student learns from its own mistakes rather than just imitating the teacher.

How much quality can you recover?

In my experience, output-based distillation from a 70B teacher to an 8B student typically recovers 85 to 95% of the teacher's quality on narrow, well-defined tasks. For broader tasks like open-ended conversation, recovery drops to 70 to 85%.

The DeepSeek-R1 distillation results are a good public reference: distilling from their 671B MoE model to Qwen-2.5 and Llama-3 bases of various sizes, they showed that a distilled 14B model could match or exceed some 70B models on math and coding benchmarks. The key was the quality and scale of the distillation data, roughly 800K reasoning trajectories.

The economics

The math usually works out strongly in favor of distillation:

• Teacher generation cost: Processing 100K prompts through a 70B model at, say, $0.50 per million input tokens and $1.50 per million output tokens costs a few hundred dollars.
• Student training cost: Fine-tuning an 8B model on 100K examples takes a few GPU-hours, roughly $10 to $50.
• Serving cost savings: If you serve 10M requests per month and the student is 4x cheaper per request, you save thousands of dollars per month.

The distillation pays for itself within the first day of serving, usually within the first hour.

Distillation is the most underused technique in production inference. If you are serving a large model and your task is well-defined, you are leaving money on the table by not distilling.

Next: building an eval harness for a deployed model, because you cannot distill, fine-tune, or quantize responsibly without knowing whether your changes helped or hurt.