Применяли ли вы дистилляцию знаний (knowledge distillation) для трансформеров?

Да, я использовал дистилляцию для уменьшения размера и ускорения больших языковых моделей (например, BERT) с минимальной потерей качества. Это критически важно для деплоя моделей в production, где есть ограничения по latency и ресурсам.

Пример реализации дистилляции для задачи классификации текста:

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 1. Загрузка предобученных моделей
teacher_name = 'bert-large-uncased'
student_name = 'bert-base-uncased'

teacher = AutoModelForSequenceClassification.from_pretrained(teacher_name, num_labels=5)
student = AutoModelForSequenceClassification.from_pretrained(student_name, num_labels=5)

# 2. Определение кастомной функции потерь для дистилляции
def distillation_loss(student_logits, teacher_logits, labels, alpha=0.5, temperature=2.0):
    """
    Комбинирует loss от hard labels и soft targets от учителя.
    """
    # Soft loss: KL-дивергенция между сглаженными распределениями
    soft_loss = nn.KLDivLoss(reduction='batchmean')(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits / temperature, dim=-1)
    ) * (temperature ** 2)

    # Hard loss: стандартный кросс-энтропийный loss для истинных меток
    hard_loss = F.cross_entropy(student_logits, labels)

    # Комбинированный loss
    return alpha * soft_loss + (1.0 - alpha) * hard_loss

# 3. Обучение студента с учетом выходов учителя (псевдокод цикла обучения)
optimizer = torch.optim.AdamW(student.parameters(), lr=5e-5)
teacher.eval()  # Учитель замораживается

for batch in dataloader:
    input_ids, attention_mask, labels = batch

    with torch.no_grad():
        teacher_logits = teacher(input_ids, attention_mask=attention_mask).logits

    student_logits = student(input_ids, attention_mask=attention_mask).logits

    loss = distillation_loss(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()

Практический результат:

• Удалось уменьшить модель для анализа тональности от bert-large (334M параметров) до distilbert-base (66M параметров). Инференс ускорился примерно в 3 раза на CPU, при падении точности (F1-score) всего на ~0.5-1%.
• Ключевые гиперпараметры: Температура (temperature) — обычно от 2 до 5. Она "смягчает" распределение вероятностей учителя, делая его более информативным для студента. Коэффициент alpha балансирует вклад soft и hard loss.