Quantization-Aware Training (QAT)

El entrenamiento consciente de la cuantización (QAT) es una técnica de entrenamiento que prepara un modelo para la inferencia de baja precisión simulando la cuantización durante el entrenamiento. Normalmente, los modelos se entrenan y almacenan con números de coma flotante de 32 bits (FP32), y la cuantización del modelo los convierte a formatos como enteros de 8 bits (INT8) para el despliegue; así, los pesos ocupan 4 veces menos y las CPU, GPU y unidades de procesamiento neuronal (NPUs) pueden realizar operaciones aritméticas enteras con rapidez. La conversión redondea cada valor y recorta los valores atípicos, lo que reduce la exactitud. QAT incorpora ese redondeo a la pasada hacia delante mientras el modelo sigue aprendiendo, para que el descenso de gradiente encuentre pesos que sigan funcionando tras la cuantización. Normalmente se ejecuta como una breve etapa de ajuste fino de un modelo ya entrenado. El esquema de inferencia basado únicamente en aritmética entera de Google para CPU móviles combinó la inferencia INT8 con este tipo de procedimiento de entrenamiento, y desde entonces QAT se ha convertido en una herramienta estándar para el despliegue de IA en el borde.

Cómo funciona el entrenamiento consciente de la cuantización#

QAT añade operaciones de cuantización simulada, también llamadas pares de cuantización y descuantización (Q/DQ), a los pesos y las activaciones del modelo. Cada operación asigna un valor flotante al nivel INT8 más cercano mediante una escala y un punto cero, y después lo convierte de nuevo inmediatamente a flotante; así, la red sigue calculando en coma flotante, pero ve los mismos valores redondeados que usaría un motor INT8. Las escalas proceden de la calibración, que registra el rango habitual de cada tensor en unos cuantos lotes de datos, aunque algunos métodos también las aprenden durante el entrenamiento.

El gradiente del redondeo es cero en casi todas partes, lo que detendría el aprendizaje. QAT utiliza el estimador de paso directo (STE), que trata el paso de redondeo como la función identidad durante la retropropagación, para que los gradientes lleguen a los pesos FP32 subyacentes:

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()}")

La pasada hacia delante usa los pesos redondeados 0.3159, −0.7247, 0.0557 y 1.18, cada uno a menos de medio paso de cuantización (0.0046) de su valor FP32; sin embargo, el gradiente que llega a los pesos FP32 es exactamente la entrada: 1.0, 2.0, 3.0 y 4.0, como si no se hubiera producido ningún redondeo. A lo largo de muchos pasos, el optimizador desplaza los pesos a posiciones en las que el error de redondeo causa el menor daño, y el modelo entrenado se exporta a INT8 con las escalas que ha aprendido.

QAT frente a la cuantización posterior al entrenamiento#

La cuantización posterior al entrenamiento (PTQ) convierte directamente un modelo FP32 ya terminado y usa unos cientos de imágenes de calibración para establecer los rangos. Tarda unos minutos y no necesita un pipeline de entrenamiento, pero los pesos no tienen ocasión de adaptarse al redondeo, por lo que los modelos sensibles pueden perder bastante exactitud. QAT requiere un ajuste fino con todos los datos de entrenamiento y, a cambio, recupera la mayor parte de esa pérdida. Los equipos suelen probar primero PTQ y recurren a QAT solo cuando la caída de exactitud en INT8 es demasiado grande:

Ultralytics YOLO ejecuta QAT cuando el entrenamiento recibe quantize=8: calibra la cuantización INT8 simulada mediante NVIDIA Model Optimizer, realiza el ajuste fino con ella en el bucle y guarda los rangos que las exportaciones a ONNX y TensorRT escriben como nodos Q/DQ sin datos de calibración aparte. En motores TensorRT sobre COCO val2017, PTQ redujo el modelo YOLO26 mediano YOLO26 de 0.527 a 0.470 de mAP_50-95, mientras que QAT lo mantuvo en 0.514; el modelo nano solo perdió 0.010 con PTQ, así que QAT no aportó nada en ese caso.

from ultralytics import YOLO

# Parte de un punto de control YOLO26 FP32 entrenado
model = YOLO("yolo26n.pt")

# Ajusta el modelo con cuantización INT8 simulada en el bucle y usa una tasa de aprendizaje baja
model.train(data="coco8.yaml", epochs=5, quantize=8, optimizer="AdamW", lr0=0.00001, cos_lr=True)

# Exporta a ONNX; los rangos INT8 aprendidos se convierten en nodos Q/DQ, por lo que no se necesitan datos de calibración
model.export(format="onnx", quantize=8)

Dónde se usa QAT#

• Visión en el borde: detectores INT8 en móviles, cámaras inteligentes, drones y NPU, donde los aceleradores que solo admiten enteros no pueden ejecutar modelos de coma flotante.
• Servidores de inferencia: motores de GPU como TensorRT ejecutan kernels Tensor Core INT8 para aumentar el rendimiento, y QAT mantiene la exactitud cerca de FP32.
• Modelos de lenguaje: QAT ya llega a menos de 8 bits. Google publicó puntos de control QAT de Gemma 3 en formato Q4_0 de 4 bits, que conservan una calidad cercana a bfloat16 con una fracción de la memoria.

Los principales frameworks incluyen herramientas de QAT, como torchao de PyTorch, el kit de herramientas de optimización de modelos de TensorFlow y NNCF de Intel para OpenVINO.

QAT frente a términos relacionados#

• Cuantización del modelo: Término general para reducir la precisión numérica. PTQ y QAT son sus dos flujos de trabajo principales.
• Entrenamiento de precisión mixta: Entrena en FP16 o BF16 para ahorrar memoria y tiempo, sin preparar el modelo para la inferencia INT8.
• Precisión media: Inferencia FP16, que la mayoría de los modelos tolera sin reentrenamiento, por lo que rara vez se necesita QAT.
• Poda de modelos: Elimina pesos o canales en lugar de reducir su precisión; ambas técnicas de compresión son compatibles.
• Destilación de conocimiento: Entrena una arquitectura más pequeña para que imite a otra más grande. A menudo se combina con QAT, con el modelo de precisión completa como profesor de su propia copia cuantizada.