Дистиляція знань (Knowledge Distillation) з великої моделі в малу

Дистиляція знань (Knowledge Distillation) з великої моделі в малу

Knowledge Distillation (KD) — техніка навчання малої моделі (student) з використанням виходів великої моделі (teacher) як «м'яких меток». Замість навчання лише на правильних відповідях (hard labels), student навчається відтворювати розподіл ймовірностей teacher'а по всьому словнику — це містить значно більше інформації про структуру задачі.

Типи дистиляції для LLM

Black-box дистиляція (Response Distillation): Використовуємо лише фінальні відповіді teacher моделі. Teacher — чорна скриня (може бути GPT-4o API). Student навчається на датасеті, де labels — виходи teacher.

# Збір даних від 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) навчається на відповідях GPT-4o через SFT

White-box дистиляція (Feature/Logit Distillation): Доступ до логітів (розподіл ймовірностей) teacher. Дозволяє навчати student на «м'яких мітках» — це більш інформативно.

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  # баланс KD та SFT loss
) -> torch.Tensor:
    """
    Комбінований loss: alpha*KD + (1-alpha)*SFT
    temperature згладжує розподіл teacher
    """
    # KD loss на м'яких мітках
    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 на 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): Замість токен-рівневих логітів, student навчається на найкращих згенерованих послідовностях teacher (beam search outputs). Простіше реалізується при black-box доступі.

DeepSeek-R1 Distill: приклад промислової дистиляції

Найвідоміший сучасний приклад — дистиляція DeepSeek-R1 (671B, MoE) в серію щільних моделей:

• DeepSeek-R1-Distill-Qwen-32B: 32B параметрів, зберігає ~85% reasoning-здібностей R1
• DeepSeek-R1-Distill-Llama-70B: 70B параметрів, ~92% від R1
• DeepSeek-R1-Distill-Llama-8B: 8B параметрів, ~70% від R1

Процес: teacher (R1) генерує 800K прикладів з CoT-міркуваннями; student навчається на них через стандартний SFT.

Практичний case study: дистиляція корпоративного асистента

Завдання: у клієнта працює GPT-4o fine-tuned для аналізу контрактів (вартість інференсу — $4000/місяць). Потрібно знизити вартість в 10× без втрати якості нижче 90% від GPT-4o рівня.

Стратегія:

1. Зібрати 12 000 запитів з продакшн-логів
2. Прогнати через GPT-4o — отримати teacher responses (датасет дистиляції)
3. Fine-tune Llama 3.1 8B на цих даних (SFT дистиляція)
4. Додатково застосувати DPO з preferred=GPT-4o відповідями, rejected=Llama baseline

Інфраструктура: збір даних через OpenAI API (~$180 на 12K запитів), навчання на A100 40GB — 6 годин.

Результати:

• Quality retention vs GPT-4o (LLM-judge): 91%
• Latency p95: 4.2с (GPT-4o API) → 0.9с (self-hosted vLLM)
• Вартість інференсу: $4000/міс → $380/міс (сервер + електрика)

Вибір temperature при дистиляції

Temperature T у KD loss визначає «м'якість» розподіли teacher:

Практика: T=3–5 для більшості задач, підбирається емпірично.

Обмеження дистиляції

• Обмеження ємності: student не може перевищити teacher, максимум наближається до teacher
• Залежність від teacher: якщо teacher помиляється, student успадковує помилки
• Вузький домен: black-box KD добре працює для спеціалізації, погано — для general capability
• Розмірний розрив: дистиляція 405B → 8B втрачає більше, ніж 70B → 8B

Часові рамки

• Збір даних від teacher: 1–3 дні
• Підготовка датасету дистиляції: 1–2 тижні
• Навчання student (8B, SFT): 3–10 годин
• Оцінка vs teacher: 3–5 днів
• Всього: 3–6 тижнів