Quantization-aware training — метод обучения, при котором модель тренируется с учётом будущего квантования, что позволяет добиться высокой точности даже при снижении разрядности весов и активаций.
Определение
Quantization-aware training (QAT) — это подход к обучению нейронных сетей, в котором операции квантования моделируются внутри тренировочного цикла. Модель обучается так, будто уже работает в низкой разрядности (например, INT8 или INT4), что позволяет минимизировать потерю качества при реальном применении квантованных весов. QAT используется, когда пост-тренировочное квантование (PTQ) приводит к слишком сильной деградации точности, особенно в больших или чувствительных архитектурах.
Как работает
Во время обучения между FP32-вычислениями вставляются симуляции квантования: округление, ограничение диапазона и обратное преобразование в FP32 для обратного прохода. Градиенты считаются в полной точности, но прямой проход учитывает эффект разрядности.
Базовый цикл QAT включает:
- имитацию квантования весов перед матричными операциями;
- имитацию квантования активаций перед передачей в следующие слои;
- STE (Straight-Through Estimator) для обратного прохода — квантование игнорируется в расчёте градиентов;
- обновление весов в FP32, но с регулярным применением ограничений, приближающих их к квантованным значениям.
Разные схемы QAT подстраивают диапазоны, масштабные коэффициенты и zero-point’ы под реальный рантайм (CPU/GPU/TPU). В LLM QAT применяют для устойчивого квантования attention-блоков, MLP и KV-кэша.
Где применяется
- LLM для инференса на CPU или мобильных устройствах.
- Компьютерное зрение (classification, detection), особенно в мобильных фреймворках.
- Edge-инференс (IoT камеры, DSP-ускорители, NPU).
- Квантованные модели для WebGPU/WebNN.
- Серверные модели, где снижение стоимости инференса критично.
Практические примеры использования
QAT широко используется в производственных пайплайнах Google TPU и NVIDIA TensorRT для задач INT8-инференса. В LLM QAT применяют при создании полностью квантованных моделей (W8A8, W4A4), включая модели для llama.cpp и ONNX Runtime. В задачах компьютерного зрения QAT является стандартом для мобильных моделей EfficientNet-Lite и MobileNet-V3, обеспечивая высокую точность на INT8-ускорителях.
Ключевые свойства
- Снижение разрывов между обучением и инференсом.
- Контроль качества при агрессивных схемах квантования.
- Совместимость с различными бэкендами (CPU, GPU, TPU, NPU).
- Возможность квантовать как веса, так и активации.
- Точная настройка scale/zero-point под рантайм.
Проблемы и ограничения
- Повышенная стоимость обучения, особенно для больших LLM.
- Сложные модели требуют отдельной настройки диапазонов для каждого слоя.
- Иногда необходимо подбирать специальные лоссы или регуляризаторы.
- Не все архитектуры устойчивы к квантованию низкой разрядности.
- Обучение дольше и требует модификаций в фреймворках.
Преимущества и ограничения
- Плюс: высокое качество квантованных моделей при минимальной деградации.
- Минус: удорожание и усложнение процесса обучения.
Связанные термины
- Post-training quantization (PTQ)
- INT8 / INT4 квантование
- Straight-through estimator
- LLM квантование
- Mixed-precision training