Дистилляция знаний (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  # баланс SFT и KD 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.

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

Задача: у клиента работает 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 ошибается, student наследует ошибки
• Узкий домен: black-box KD хорошо работает для specialization, плохо — для general capability
• Размерный разрыв: дистилляция 405B → 8B теряет больше, чем 70B → 8B

Сроки

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