Оптимизация нейросетей: квантование и дистилляция

Кратко

Квантование и дистилляция — два основных способа сжать обученную нейросеть, чтобы она работала быстрее и помещалась на слабом железе. Квантование понижает точность чисел (обычно int8 вместо 32-битных float), дистилляция переносит «знание» большой модели в маленькую. Без этих приёмов современные модели не влезали бы ни в смартфон, ни в бортовой компьютер робота.

Для чего используется

запуск моделей на встраиваемых и мобильных устройствах: меньше памяти, ниже задержка и энергопотребление;
совместимость со специализированными ускорителями: например, Google Edge TPU исполняет только полностью квантованные (int8) модели;
удешевление инференса в облаке: та же нагрузка на меньшем числе серверов;
уменьшение размера дистрибутива: меньше места на устройстве и трафика на загрузку.

Как работает

Квантование (quantization) — вычисления и хранение данных с пониженной точностью: обычно int8 вместо fp32. Выгоды по данным команды PyTorch: модель в 4 раза меньше, потребность в пропускной способности памяти ниже в 2–4 раза, инференс быстрее в 2–4 раза (точное ускорение зависит от железа, рантайма и модели). Аналогия: перейти от измерений с точностью до миллиметра к сантиметрам — записи короче, работа быстрее, а для большинства задач точности хватает.

Три подхода:

динамическое — веса переводятся в int8 заранее, активации квантуются на лету; подходит трансформерам и LSTM;
статическое post-training — диапазоны активаций калибруются прогоном небольшого набора данных через «наблюдатели»; подходит свёрточным сетям;
quantization-aware training (QAT) — обучение с имитацией квантования (float-значения округляются, как будто они int8); применяют, когда post-training теряет слишком много точности.

Цифры Google LiteRT: post-training float16 сокращает модель до 50%, dynamic range и полное int8 — до 75%. Пример MobileNet-v1: размер 16,9 → 4,3 МБ (−75%), задержка 124 → 64 мс при QAT.

Дистилляция знаний (knowledge distillation) — обучение маленькой модели-ученика воспроизводить поведение большой модели-учителя (Hinton, Vinyals, Dean; 2015). Ученик учится не на жёстких метках, а на «мягких» вероятностях учителя (температурный softmax) — они несут больше информации о том, как учитель различает классы. Мотивация прямая: предсказывать целым ансамблем «громоздко и слишком дорого для развёртывания», а его знание можно сжать в одну компактную модель. Метод улучшил акустическую модель коммерческой системы Google.

Практический эталон — DistilBERT (Hugging Face, 2019): модель на 40% меньше BERT, сохраняет 97% качества понимания языка и работает на 60% быстрее; дистилляция применена уже на этапе предобучения.

Смежные приёмы. Прунинг (pruning) удаляет параметры с малым вкладом; кластеризация весов группирует их вокруг общих центроидов, сокращая число уникальных значений. Оба улучшают сжимаемость модели.

Основные компоненты

форматы чисел — fp32 → fp16 → int8 (и ниже);
калибровка — подбор диапазонов квантования по реальным данным;
учитель и ученик — пара моделей в дистилляции, «мягкие» вероятности;
инструменты — PyTorch quantization, Google LiteRT (бывш. TensorFlow Lite).

Примеры

MobileNet-v1 + QAT (LiteRT) — 16,9 → 4,3 МБ, 124 → 64 мс.
DistilBERT — −40% размера, 97% качества, +60% скорости относительно BERT.

Связанные темы

Источники

Introduction to Quantization on PyTorch — определение, выгоды, три подхода.
Google LiteRT — Model optimization — цифры сжатия, пример MobileNet, прунинг и кластеризация, требования Edge TPU.
Distilling the Knowledge in a Neural Network (Hinton et al., 2015) — идея и мотивация дистилляции.
DistilBERT (Sanh et al., 2019) — практический пример дистилляции.