Knowledge distillation

Knowledge distillation — это метод машинного обучения, в котором компактная модель (ученик) обучается воспроизводить поведение более крупной, сложной модели (учителя) с целью достижения сопоставимой точности при меньших вычислительных затратах и объёме памяти. В отличие от традиционного обучения, где модель учится на «жёстких» метках (например, «кошка» или «собака»), дистилляция знаний использует «мягкие» предсказания учителя — распределение вероятностей по всем классам, которое содержит дополнительную информацию о сходстве между классами (например, что изображение «кошки» немного похоже на «рысь», но не на «автомобиль»). Термин был введён Джеффри Хинтоном, Ориолем Виньялсом и Джеффом Дином в 2015 году.

История

Идея передачи знаний от одной модели к другой восходит к началу 1990-х годов, когда исследователи изучали методы сжатия нейронных сетей. В 1991 году Юрген Шмидхубер предложил концепцию «сжатия нейронных сетей», где одна сеть обучается имитировать выходы другой. Однако практический прорыв произошёл в 2015 году, когда группа Хинтона опубликовала статью «Distilling the Knowledge in a Neural Network». Авторы показали, что использование «мягких» меток с температурным параметром позволяет ученику улавливать тонкие взаимосвязи между классами, которые не видны в «жёстких» метках. С тех пор метод получил широкое распространение в компьютерном зрении, обработке естественного языка и рекомендательных системах.

Принцип работы

Архитектура учитель-ученик

Процесс дистилляции знаний предполагает наличие двух моделей:

• Модель-учитель — большая, глубокая или ансамблевая модель, обученная на исходных данных. Она может быть слишком ресурсоёмкой для развёртывания на устройствах с ограниченными вычислительными возможностями (например, мобильных телефонах или встраиваемых системах).
• Модель-ученик — компактная модель с меньшим числом параметров, которая обучается подражать учителю.

Мягкие метки и температура

Ключевой элемент дистилляции — преобразование выходных вероятностей учителя с помощью параметра температуры \( T \). Формула softmax с температурой выглядит так:

\[ q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} \]

где \( z_i \) — логиты (выходы последнего слоя) для класса \( i \), а \( T \) — температура. При \( T = 1 \) получается стандартный softmax. При \( T > 1 \) распределение становится более «мягким» — вероятности выравниваются, что подчёркивает относительные сходства между классами. При \( T < 1 \) распределение становится более «жёстким», приближаясь к one-hot кодированию.

Функция потерь

Обучение ученика происходит путём минимизации комбинированной функции потерь, которая включает два компонента:

1. Дистилляционная потеря — расхождение Кульбака-Лейблера (KL-дивергенция) между мягкими предсказаниями учителя и мягкими предсказаниями ученика (с одинаковой температурой \( T \)).
2. Стандартная потеря — кросс-энтропия между предсказаниями ученика и истинными метками (обычно с температурой \( T = 1 \)).

Итоговая функция потерь имеет вид:

\[ L = \alpha \cdot L_{\text{soft}} + (1 - \alpha) \cdot L_{\text{hard}} \]

где \( \alpha \) — гиперпараметр, регулирующий баланс между двумя компонентами. На практике часто используют \( T = 2...5 \) и \( \alpha = 0.7...0.9 \).

Виды дистилляции знаний

По способу передачи знаний

• Offline дистилляция — учитель предварительно обучен и заморожен; ученик обучается на его предсказаниях. Наиболее распространённый вариант.
• Online дистилляция — учитель и ученик обучаются одновременно; учитель может быть той же архитектуры, что и ученик, или динамически обновляться.
• Self-distillation — ученик обучается на собственных предсказаниях, полученных на предыдущих итерациях (например, в подходе «Born-Again Networks»).

По типу знаний

• Response-based — ученик имитирует выходы учителя (логиты или вероятности).
• Feature-based — ученик учится воспроизводить внутренние представления (карты признаков) из промежуточных слоёв учителя.
• Relation-based — ученик обучается воспроизводить отношения между образцами (например, матрицы попарных сходств).

По архитектуре

• Гомогенная дистилляция — учитель и ученик имеют одинаковую архитектуру, но разный размер (например, ResNet-50 учит ResNet-18).
• Гетерогенная дистилляция — архитектуры различаются (например, трансформер учит свёрточную сеть).

Применение

Сжатие моделей

Основное применение дистилляции знаний — уменьшение размера модели без существенной потери точности. Например, BERT-base (110 млн параметров) может быть дистиллирован в DistilBERT (66 млн параметров), который сохраняет 95% производительности при 40% меньшем размере и 60% более быстром выводе.

Мобильные и встраиваемые системы

Дистилляция позволяет развёртывать нейросетевые модели на устройствах с ограниченными ресурсами — смартфонах, камерах, IoT-датчиках. Например, модель MobileNet может быть обучена как ученик более крупной ResNet-50.

Обучение с учителем без меток

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

Ансамблирование

Дистилляция позволяет сжать ансамбль моделей в одну компактную модель. Например, ансамбль из 10 ResNet-50 может быть дистиллирован в одну ResNet-50, которая превосходит по точности каждую отдельную модель из ансамбля.

Преимущества и недостатки

Преимущества

• Эффективность — компактные модели требуют меньше памяти и вычислительных ресурсов.
• Сохранение точности — при правильной настройке ученик может достигать точности, близкой к учителю.
• Гибкость — метод не зависит от конкретной архитектуры и может применяться к различным типам моделей.
• Улучшение обобщения — мягкие метки учителя содержат регуляризующую информацию, снижающую переобучение.

Недостатки

• Зависимость от учителя — качество ученика напрямую зависит от качества учителя; плохой учитель может передать ученику свои ошибки.
• Вычислительные затраты на обучение — процесс требует обучения учителя (часто ресурсоёмкого) и дополнительного этапа дистилляции.
• Чувствительность к гиперпараметрам — температура \( T \) и коэффициент \( \alpha \) требуют тщательной настройки.
• Ограниченная применимость — для некоторых задач (например, регрессии с непрерывным выходом) дистилляция менее эффективна.

Связь с другими методами

Дистилляция знаний тесно связана с такими техниками, как:

• Прунинг — удаление избыточных параметров модели; часто используется вместе с дистилляцией.
• Квантование — снижение точности чисел (например, с 32-битных float до 8-битных int); может применяться после дистилляции.
• Архитектурный поиск — автоматический поиск компактных архитектур; дистилляция может служить критерием оценки.

В отличие от трансферного обучения, где знания переносятся между задачами, дистилляция фокусируется на переносе знаний между моделями одной задачи.

Критика

Основная критика дистилляции знаний связана с тем, что она не всегда даёт ожидаемый выигрыш. В ряде исследований (например, работа «Do We Really Need to Access the Source Data?» 2020 года) показано, что при обучении ученика на «жёстких» метках с использованием аугментации данных можно достичь результатов, сопоставимых с дистилляцией. Кроме того, метод требует доступа к предсказаниям учителя, что может быть проблематично в условиях конфиденциальности данных (например, в медицинских приложениях). Некоторые авторы отмечают, что дистилляция может приводить к «катастрофическому забыванию» знаний, если учитель и ученик имеют сильно различающиеся архитектуры.

Источники

• Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv preprint arXiv:1503.02531.
• Buciluă, C., Caruana, R., & Niculescu-Mizil, A. (2006). Model Compression. Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining.
• Sanh, V., Debut, L., Chaumond, J., & Wolf, T. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108.
• Gou, J., Yu, B., Maybank, S. J., & Tao, D. (2021). Knowledge Distillation: A Survey. International Journal of Computer Vision, 129(6), 1789–1819.