Knowledge Distillation

Дистилляция знаний — это метод сжатия моделей, при котором небольшую «ученическую» нейронную сеть обучают воспроизводить поведение более крупной и точной модели-«учителя». Ученик обучается на выходных вероятностях или внутренних признаках учителя в дополнение к эталонным меткам либо вместо них. Эти более информативные целевые значения позволяют добиться большей точности, чем при обучении той же архитектуры только на метках. В результате получается модель, достаточно быстрая и компактная для телефонов, встроенных плат и другого оборудования для периферийного ИИ, которая сохраняет значительную часть качества модели-учителя. В 2006 году Буцилуэ, Карауна и Никулеску-Мизил показали, что знания ансамбля можно сжать в одну модель, а Хинтон, Виньялс и Дин в 2015 году назвали этот подход дистилляцией знаний и представили смягчённые температурой целевые значения, которые до сих пор используются в большинстве методов.

Как работает дистилляция знаний#

Метка сообщает только, что на изображении изображён волк. Выход softmax обученного учителя сообщает больше: изображение немного похоже на хаски, слегка — на кошку и совсем не похоже на машину. Эти относительные вероятности неправильных классов, часто называемые тёмными знаниями, передаются ученику в виде мягких целевых значений. Деление логитов на температуру T, превышающую 1, перед применением softmax сглаживает распределение, чтобы небольшие вероятности тоже вносили достаточный вклад в обучение. Функция потерь дистилляции — это дивергенция Кульбака — Лейблера между сглаженными распределениями учителя и ученика, умноженная на T², чтобы градиенты оставались сопоставимыми по величине с градиентами обычной функции потерь по меткам:

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

При T=1 учитель присваивает классу «волк» вероятность 0.875, а классам «кошка» и «машина» — всего 0.006 и 0.0. При T=4 те же логиты дают значения 0.493, 0.299, 0.141 и 0.067, показывая ученику, что волк больше похож на хаски, чем на кошку, а кошка — больше на волка, чем на машину. Ученик минимизирует эту функцию потерь, здесь равную 1.244, обычно в виде взвешенной суммы с обычной функцией потерь задачи по эталонным меткам. Учитель остаётся замороженным на протяжении всего процесса и выполняет только прямые проходы:

Виды дистилляции#

• На основе выходов: ученик подгоняет свои окончательные выходы под выходы учителя, как в описанном выше методе мягких целевых значений. От учителя требуются только логиты. DistilBERT объединил этот подход с языковым моделированием и функциями потерь по скрытым состояниям во время предварительного обучения. В результате BERT стал на 40% меньше и на 60% быстрее, сохранив 97% способности понимать язык.
• На основе признаков: ученик подгоняет промежуточные карты признаков учителя, а небольшая обучаемая проекционная сеть компенсирует возможную разницу в числе каналов. В работе FitNets были представлены такие «подсказки» и обучен ученик примерно с в 10.4 раза меньшим числом параметров, чем у учителя, который превзошёл учителя на CIFAR-10.
• На уровне последовательностей: учитель — языковая модель — генерирует текст, а ученик дообучается на этих выходах. DeepSeek дообучил свои дистиллированные модели R1, основанные на Qwen2.5 и Llama 3, на 800,000 примерах, отобранных с помощью DeepSeek-R1.

Дистилляция на основе признаков подходит для детекторов, выходы которых содержат тысячи рамок и оценок, а не одно распределение классов. Ultralytics YOLO использует этот подход, если при обучении передан учитель distill_model: замороженная более крупная модель YOLO26 предоставляет три карты признаков шеи, которые поступают в её голову детекции, а ученик подгоняет свои карты под них с помощью проекционных свёрток 1×1 и функции потерь L2, взвешенной уверенностью учителя и масштабированной параметром dis (по умолчанию 6.0). На COCO val2017 дистиллированная нано-модель YOLO26 достигает 41.5 mAP50-95 против 40.9 у такой же модели, обученной без учителя, при одинаковой стоимости инференса.

from ultralytics import YOLO

# Ученик: компактная нано-модель YOLO26, предназначенная для развёртывания
model = YOLO("yolo26n.pt")

# Учитель: более крупный чекпойнт YOLO26, который направляет обучение и затем удаляется
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

Дистилляция и родственные термины#

• Трансферное обучение: использует собственные веса модели как начальную точку для новой задачи. При дистилляции другая, обычно меньшая архитектура обучается имитировать модель для той же задачи.
• Обрезка модели: удаляет веса, каналы или слои из обученной сети. Чтобы восстановить точность, обрезанные модели часто дообучают с помощью дистилляции исходной модели.
• Квантование модели: хранит веса и активации с меньшей точностью, например INT8. Оно уменьшает ту же архитектуру и сочетается с дистилляцией, которая позволяет выбрать архитектуру меньшего размера.
• Ансамбль моделей: усредняет результаты нескольких моделей во время инференса. Изначальной мотивацией этого метода была дистилляция ансамбля в одну модель-ученика.
• Дистилляция набора данных: несмотря на сходство названий, она сжимает набор данных, а не модель. Ван и соавторы сжали 60,000 обучающих изображений MNIST до 10 синтетических изображений — по одному на класс. С ними сеть с фиксированной инициализацией за несколько шагов градиентного спуска достигает точности, близкой к исходной.