Quantization-Aware Training (QAT)

Обучение с учётом квантования (QAT) — это метод обучения, который подготавливает модель к инференсу с низкой точностью, имитируя квантование в процессе обучения. Обычно модели обучают и хранят веса в формате 32-разрядных чисел с плавающей точкой (FP32), а квантование модели преобразует их, например, в 8-разрядные целые числа (INT8) для развёртывания. Это уменьшает размер весов в 4 раза и позволяет CPU, GPU и нейропроцессорам (NPU) выполнять быстрые целочисленные вычисления. При преобразовании значения округляются, а выбросы отсекаются, что снижает точность. QAT добавляет такое округление в прямой проход, пока модель ещё обучается, чтобы градиентный спуск находил веса, которые сохраняют работоспособность после квантования. Обычно QAT применяют на коротком этапе дообучения уже обученной модели. В разработанной Google схеме инференса только с целочисленной арифметикой для мобильных CPU инференс INT8 сочетался с таким способом обучения. С тех пор QAT стал стандартным инструментом для развёртывания периферийного ИИ.

Как работает обучение с учётом квантования#

QAT добавляет к весам и активациям модели операции фиктивного квантования, также называемые парами «квантование-деквантование» (Q/DQ). Каждая такая операция преобразует число с плавающей точкой в ближайшее значение INT8 с помощью масштаба и нулевой точки, а затем сразу преобразует его обратно. Поэтому вычисления в сети по-прежнему выполняются с числами с плавающей точкой, но используются те же округлённые значения, что и в движке INT8. Масштабы определяются калибровкой: она фиксирует типичный диапазон каждого тензора на нескольких пакетах данных. Некоторые методы также обучают эти значения.

Почти везде градиент округления равен нулю, что остановило бы обучение. QAT использует оцениватель со сквозным распространением ошибки (STE): при обратном распространении ошибки округление считается тождественной функцией, поэтому градиенты проходят к исходным весам FP32:

import torch

# FP32 weights that training keeps and updates
w = torch.tensor([0.3137, -0.7214, 0.0521, 1.1800], requires_grad=True)

# Symmetric INT8 fake quantization: snap each weight to one of 255 levels, then map back to float
scale = w.detach().abs().max().item() / 127
w_q = torch.fake_quantize_per_tensor_affine(w, scale, 0, -127, 127)
print(f"Fake-quantized: {[round(v, 4) for v in w_q.tolist()]}")

# The forward pass sees the rounded weights; the straight-through estimator still delivers gradients to w
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
loss = (w_q * x).sum()
loss.backward()
print(f"Gradient on FP32 weights: {w.grad.tolist()}")

В прямом проходе используются округлённые веса 0.3159, −0.7247, 0.0557 и 1.18; каждый отличается от значения FP32 не более чем на половину шага квантования (0.0046). При этом градиент, достигающий весов FP32, в точности равен входным значениям: 1.0, 2.0, 3.0 и 4.0, как если бы округления не было. За много шагов оптимизатор перемещает веса в такие позиции, где ошибка округления наносит наименьший ущерб, а обученная модель экспортируется в INT8 с выученными масштабами.

QAT и квантование после обучения#

Квантование после обучения (PTQ) напрямую преобразует готовую модель FP32, используя несколько сотен калибровочных изображений для определения диапазонов. Процесс занимает минуты и не требует конвейера обучения, но веса не могут адаптироваться к округлению, поэтому у чувствительных моделей заметно снижается точность. Для QAT требуется дообучение на полном наборе данных, зато он восстанавливает большую часть потерянной точности. Обычно команды сначала пробуют PTQ и переходят на QAT, только если падение точности INT8 слишком велико:

Ultralytics YOLO выполняет QAT, если при обучении передан параметр quantize=8: калибровка фиктивного квантования INT8 выполняется через NVIDIA Model Optimizer, затем модель дообучается с включённым квантованием, а сохранённые диапазоны при экспорте в ONNX и TensorRT записываются как узлы Q/DQ без отдельной калибровки. В движках TensorRT на COCO val2017 PTQ снизил показатель mAP_50-95 средней модели YOLO26 с 0.527 до 0.470, тогда как с QAT он составил 0.514. У нано-модели PTQ снизил показатель всего на 0.010, поэтому QAT ничего не добавил.

from ultralytics import YOLO

# Начни с обученной контрольной точки FP32 YOLO26
model = YOLO("yolo26n.pt")

# Выполни дообучение с фиктивным квантованием INT8 в цикле и используй небольшую скорость обучения
model.train(data="coco8.yaml", epochs=5, quantize=8, optimizer="AdamW", lr0=0.00001, cos_lr=True)

# Экспортируй модель в ONNX; выученные диапазоны INT8 станут узлами Q/DQ, поэтому калибровочные данные не понадобятся
model.export(format="onnx", quantize=8)

Где применяют QAT#

• Компьютерное зрение на периферии: детекторы INT8 для мобильных телефонов, умных камер, дронов и NPU, где ускорители, работающие только с целыми числами, вообще не могут запускать модели с плавающей точкой.
• Серверы инференса: графические движки, такие как TensorRT, запускают ядра Tensor Core INT8 для повышения пропускной способности, а QAT сохраняет точность на уровне, близком к FP32.
• Языковые модели: QAT теперь позволяет использовать менее 8 бит. Google выпустила контрольные точки Gemma 3 QAT в 4-разрядном формате Q4_0, которые сохраняют качество, близкое к bfloat16, при значительно меньшем объёме памяти.

Крупные фреймворки включают инструменты QAT: torchao для PyTorch, Model Optimization Toolkit для TensorFlow и NNCF от Intel для OpenVINO.

QAT и связанные с ним термины#

• Квантование модели: общий термин для снижения числовой точности. PTQ и QAT — два основных способа его выполнения.
• Обучение с пониженной точностью: обучение в формате FP16 или BF16 для экономии памяти и времени без подготовки модели к инференсу INT8.
• Половинная точность: инференс FP16, который большинство моделей поддерживают без переобучения, поэтому QAT для него требуется редко.
• Прореживание модели: удаление весов или каналов вместо снижения их точности; эти способы сжатия можно сочетать.
• Дистилляция знаний: обучение архитектуры меньшего размера для имитации более крупной. Её часто сочетают с QAT: полноточная модель выступает учителем для собственной квантованной копии.