Оптимизация ML-модели (дистилляция) для мобильного устройства

Мобильная нейросеть должна быть лёгкой, но точной. Knowledge Distillation — метод, который позволяет сжать большую модель без критической потери качества. Мы применили этот подход на десятках проектов: от классификации изображений до NLP на устройстве. Результат — модель в 10 раз меньше при падении точности всего на 1–2%. Такой подход особенно востребован в мобильной ML, где каждый мегабайт и миллисекунда на счету. Свяжитесь с нами — мы проведём аудит вашей модели и предложим план дистилляции.

Почему мягкие метки работают лучше

Обычное обучение: правильный класс = 1.0, остальные = 0.0. Hard labels. Учитель на изображении кошки выдаёт: кошка 0.85, рысь 0.08, тигр 0.04, собака 0.02 ... Эти «мягкие» метки несут информацию, что рысь похожа на кошку больше, чем самолёт. Студент, обученный на таких метках, усваивает структуру пространства признаков, а не просто решение бинарного классификатора. Это суть парадигмы student-teacher.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, true_labels, temperature=4.0, alpha=0.7):
    """ alpha — вес дистилляции vs hard label loss
    temperature — сглаживает распределение учителя
    """
    # Soft targets loss (KL divergence между студентом и учителем)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    distill_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (temperature ** 2)

    # Hard label loss (обычная кросс-энтропия)
    hard_loss = F.cross_entropy(student_logits, true_labels)

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

temperature ** 2 — нормализующий множитель, компенсирующий масштаб градиентов при высокой температуре. Без него distill_loss и hard_loss находятся в разных масштабах.

Выбор архитектуры студента

Студент должен быть меньше учителя, но не произвольно. Хорошие базовые архитектуры для мобиля:

• MobileNetV3-Small — 2.5 МБ, проектировался для мобиля с нуля, depthwise separable convolutions
• EfficientNet-Lite0/1 — хороший баланс точность/скорость
• MobileViT-XXS — hybrid CNN+Transformer, 1.3 МБ
• DistilBERT (для NLP) — уже дистиллированный из BERT, 66 МБ vs 440 МБ

Для задач детекции объектов на мобиле: студент на базе YOLOv8n (8 МБ) дистиллируется из YOLOv8l (87 МБ).

Как сравниваются методы дистилляции?

Существует несколько подходов к дистилляции, различающихся по сложности и результату. Основные: логит-дистилляция (по выходам), дистилляция промежуточных слоёв и data-free дистилляция. Выбор зависит от доступности данных и требований к точности.

Этапы дистилляции (шаги)

1. Подготовка учителя: дообучить или квантизировать существующую модель, заморозить веса.
2. Выбор архитектуры студента: под ресурсный бюджет (RAM, CPU, latency).
3. Настройка гиперпараметров: temperature (обычно 2-8), alpha (0.5-0.9), выбор слоёв для промежуточной дистилляции.
4. Обучение студента: минимизация комбинированного loss, мониторинг валидационной точности.
5. Конвертация в Core ML/TFLite с квантизацией int8 для финального размера.
6. Валидация на реальных устройствах: проверка латенси и точности.

Процесс дистилляции: пример для классификации

# Предположим: учитель — ResNet-50, студент — MobileNetV3-Small
teacher = torchvision.models.resnet50(pretrained=True).eval()
student = torchvision.models.mobilenet_v3_small(pretrained=False)

# Замораживаем учителя
for param in teacher.parameters():
    param.requires_grad = False

optimizer = torch.optim.AdamW(student.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(100):
    student.train()
    for images, labels in train_loader:
        with torch.no_grad():
            teacher_logits = teacher(images)
        student_logits = student(images)
        loss = distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()
    val_acc = evaluate(student, val_loader)
    print(f"Epoch {epoch}: student_acc={val_acc:.4f}")

Типичные результаты: MobileNetV3-Small обученный обычно — 67–68% top-1 на ImageNet. После дистилляции из ResNet-50 — 71–72%. Прирост 3–4% за счёт knowledge transfer.

Как работает intermediate layer distillation?

Дистилляция только по выходам (logits) — базовый вариант. Более сильный: добавляем соответствие промежуточных feature maps.

# FitNets / PKT: студент учит feature maps учителя
class DistillationHook:
    """Хук для захвата промежуточных активаций"""
    def __init__(self):
        self.output = None
    def __call__(self, module, input, output):
        self.output = output

teacher_hook = DistillationHook()
student_hook = DistillationHook()

# Регистрируем на соответствующих слоях
teacher.layer3.register_forward_hook(teacher_hook)
student.features[9].register_forward_hook(student_hook)  # Analogous layer

# В цикле обучения добавляем feature distillation loss
with torch.no_grad():
    teacher(images)
    teacher_features = teacher_hook.output

student(images)  # с grad
student_features = student_hook.output

# Если размерности отличаются — нужен adapter (1x1 Conv)
if teacher_features.shape[1] != student_features.shape[1]:
    student_features = adapter_conv(student_features)  # adapter обучается вместе

feature_loss = F.mse_loss(student_features, teacher_features.detach())

Такой подход требует выравнивания размерностей feature maps между учителем и студентом — через adapter 1×1 свёрток. Адаптер добавляет немного параметров студенту, но остаётся маленьким.

Как дистиллировать модель, если исходные данные недоступны?

Иногда исходный датасет недоступен (IP restrictions, privacy). Data-free distillation — генерируем синтетические данные, которые максимизируют активации учителя:

# DAFL (Data-Free Learning): генератор создаёт «образцы» для дистилляции
generator = Generator(latent_dim=256, img_channels=3)
optimizer_G = torch.optim.Adam(generator.parameters(), lr=1e-4)
for step in range(1000):
    z = torch.randn(batch_size, 256)
    fake_images = generator(z)
    # Потери: максимизируем уверенность учителя + минимизируем BatchNorm statistics mismatch
    teacher_out = teacher(fake_images)
    activation_loss = -teacher_out.max(dim=1)[0].mean()  # учитель должен быть уверен
    # BN statistics matching
    bn_loss = compute_bn_statistics_loss(teacher, fake_images)
    total_loss = activation_loss + 0.1 * bn_loss
    optimizer_G.zero_grad()
    total_loss.backward()
    optimizer_G.step()

Качество data-free дистилляции ниже полноданных варианта, но иногда это единственный вариант.

Дистилляция для NLP задач на мобиле

Для мобильных приложений с NLP (классификация отзывов, определение intent, суммаризация): дистиллируем из GPT-4 / Claude API ответов в маленький BERT/DistilBERT.

# Собираем soft labels от учителя (GPT-4 API)
# Для каждого обучающего примера запрашиваем вероятности классов
# Сохраняем как обучающие метки для студента
# Студент — DistilBERT fine-tuned на этих мягких метках

DistilBERT (66 МБ, ONNX int8 — 18 МБ) работает на устройстве за 30–80 мс на iOS/Android. GPT-4 в облаке — сотни мс, деньги за запросы.

Что входит в работу

• Анализ ресурсного бюджета: RAM, CPU, GPU, энергопотребление
• Выбор архитектуры студента под ваши ограничения
• Подготовка учителя (дообучение, квантизация)
• Дистилляция с подбором гиперпараметров (temperature, alpha, intermediate layers)
• Валидация на тестовой выборке и на реальных устройствах
• Конвертация в Core ML / TFLite с квантизацией int8
• Интеграция и сопровождение (апдейты при новых версиях ОС)

Наша команда имеет 10+ лет опыта в мобильной разработке и ML. Мы гарантируем, что модель будет стабильно работать на устройствах с iOS 14+ и Android 10+. Получите консультацию инженера по вашей модели. Закажите дистилляцию — мы подготовим предварительную оценку за 1 день.

Ориентиры по срокам

Базовая логит-дистилляция для классификационной задачи — 2–4 недели (GPU-время плюс подбор гиперпараметров). Полная дистилляция с промежуточными слоями, нестандартными архитектурами, data augmentation — 5–10 недель. Индивидуальный расчёт — по запросу.

Дополнительно: Knowledge Distillation на Wikipedia, туториал PyTorch.