Как дистиллировать одну модель в другую?

Дистилляция модели — это процесс переноса знаний из большой, сложной модели (teacher) в более компактную и быструю модель (student). Цель — сохранить качество предсказаний при уменьшении размера и вычислительной нагрузки.

Основные шаги:

1. Обучить teacher-модель на исходных данных.
2. Использовать teacher для генерации «мягких» меток (soft labels) — вероятностных распределений по классам, а не только жёстких ответов.
3. Обучить student-модель, минимизируя комбинированную функцию потерь, которая учитывает:
   • Расхождение между предсказаниями student и soft labels teacher (обычно с помощью кросс-энтропии с температурой).
   • Расхождение между предсказаниями student и истинными метками (если они есть).

Пример на PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

# teacher и student — модели
# inputs, labels — данные

temperature = 2.0
alpha = 0.7  # вес для soft loss

# Получаем предсказания teacher с температурой
with torch.no_grad():
    teacher_logits = teacher(inputs) / temperature
    soft_targets = F.softmax(teacher_logits, dim=1)

student_logits = student(inputs) / temperature

# Потеря для soft targets
loss_soft = F.kl_div(F.log_softmax(student_logits, dim=1), soft_targets, reduction='batchmean') * (temperature ** 2)

# Потеря для истинных меток
loss_hard = F.cross_entropy(student(inputs), labels)

loss = alpha * loss_soft + (1 - alpha) * loss_hard

loss.backward()

Дистилляция помогает ускорить inference и уменьшить требования к памяти, сохраняя при этом приемлемое качество модели.