Knowledge Distillation from Large to Small Model

Knowledge Distillation from large to small model

We apply Knowledge Distillation (KD) to reduce LLM inference costs without significant quality loss. KD trains a small model (student) using outputs of a large model (teacher) as soft labels. The student learns to reproduce the teacher's probability distribution across the vocabulary. This carries far more information than a single correct answer. Our team has completed 30+ distillation projects — from QA systems to contract analysis assistants. In practice, distillation preserves 85–95% of teacher quality at a fraction of the inference cost. We work with models from 7B to 405B parameters. If you run a large LLM in production and want to reduce costs 5–10×, contact us to assess your project. We will evaluate your current model and propose the optimal compression strategy.

Types of distillation for LLMs

Black-box distillation (Response Distillation): Use only final answers from teacher model. Teacher is a black box (can be GPT-4o API). Student is trained on a dataset where labels are teacher outputs.

# Collect data from teacher (GPT-4o)
def collect_teacher_outputs(prompts: list[str], client) -> list[dict]:
    dataset = []
    for prompt in prompts:
        teacher_response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        ).choices[0].message.content

        dataset.append({"prompt": prompt, "response": teacher_response})
    return dataset

# Student (Llama 3.1 8B) trained on GPT-4o answers via SFT

White-box distillation (Feature/Logit Distillation): Access to teacher logits (probability distribution). Allows training student on "soft labels" — more informative.

import torch
import torch.nn.functional as F

def distillation_loss(
    student_logits,    # [batch, seq_len, vocab_size]
    teacher_logits,    # [batch, seq_len, vocab_size]
    labels,            # [batch, seq_len]
    temperature: float = 4.0,
    alpha: float = 0.5  # balance KD and SFT loss
) -> torch.Tensor:
    """
    Combined loss: alpha*KD + (1-alpha)*SFT
    temperature smooths teacher distribution
    """
    # KD loss on soft labels
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    kd_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)

    # SFT loss on hard labels
    sft_loss = F.cross_entropy(
        student_logits.view(-1, student_logits.size(-1)),
        labels.view(-1),
        ignore_index=-100
    )

    return alpha * kd_loss + (1 - alpha) * sft_loss

Sequence-level KD (SeqKD): Instead of token-level logits, student trains on best generated sequences from teacher (beam search outputs). Simpler to implement with black-box access.

Which distillation method to choose?

DeepSeek-R1 Distill: example of industrial distillation

Most known modern example — distillation of DeepSeek-R1 (671B, MoE) into series of dense models:

• DeepSeek-R1-Distill-Qwen-32B: 32B parameters, retains ~85% of R1 reasoning ability.
• DeepSeek-R1-Distill-Llama-70B: 70B parameters, ~92% of R1.
• DeepSeek-R1-Distill-Llama-8B: 8B parameters, ~70% of R1.

Process: teacher (R1) generates 800K examples with CoT reasoning. Student trains on them via standard SFT.

Practical case study: corporate assistant distillation

Task: our client ran a fine-tuned GPT-4o for contract analysis. Inference costs were substantial. The goal was to reduce costs 10× without quality dropping below 90% of GPT-4o level.

Strategy applied in this project:

1. Collect 12,000 requests from production logs.
2. Run through GPT-4o to get teacher responses (distillation dataset).
3. Fine-tune Llama 3.1 8B on this data via SFT distillation.
4. Apply DPO with preferred=GPT-4o answers and rejected=Llama baseline.

Infrastructure: data collection via OpenAI API, training on A100 40GB — 6 hours. Data collection cost is recovered in the first week of operation.

Results from our practice:

• Quality retention vs GPT-4o (LLM-judge): 91%.
• Latency p95: 4.2s (GPT-4o API) → 0.9s (self-hosted vLLM).
• Inference cost: reduced by more than 10× on self-hosted vLLM.

Temperature selection in distillation

Temperature T in KD loss determines "softness" of teacher distribution:

Practice: T=3–5 for most tasks, selected empirically.

Distillation limitations

• Capacity bottleneck: student cannot exceed teacher; maximum approaches teacher level.
• Teacher dependency: if teacher makes mistakes, student inherits them.
• Narrow domain: black-box KD works well for specialization, poorly for general capability.
• Size gap: distilling 405B → 8B loses more than 70B → 8B.

What Is Included in Our Distillation Service

Our knowledge distillation service delivers the full pipeline:

• Analysis of your current model and target quality metrics.
• Distillation dataset collection and preparation from teacher or from production logs.
• Student model training: architecture selection and hyperparameter tuning.
• Evaluation against teacher: LLM-judge, accuracy metrics, latency p99.
• Inference optimization: quantization, vLLM, ONNX Runtime.
• Documentation and team training included.
• Technical support during launch.

We guarantee that the final model retains at least 90% of teacher quality on key metrics. Contact us to get an accurate assessment and timeline for your project.

Timeline

• Collecting data from teacher: 1–3 days.
• Distillation dataset preparation: 1–2 weeks.
• Student training (8B, SFT): 3–10 hours.
• Evaluation vs teacher: 3–5 days.
• Total project duration: 3–6 weeks.