Optimización de XNNPack para inferencia cuantizada

Google ha implementado una actualización técnica en XNNPack, el backend de CPU de TensorFlow Lite, que permite la ejecución de inferencias mediante cuantización de rango dinámico en los operadores de Convolución 2D y Fully Connected. Esta optimización busca maximizar el rendimiento de la inteligencia artificial en el hardware más extendido, ya que las CPUs siguen siendo el objetivo predeterminado para la ejecución de modelos en dispositivos móviles y embebidos.

La cuantización es un proceso que reduce la precisión de los pesos de un modelo, pasando de números de coma flotante de 32 bits (fp32) a enteros de 8 bits (int8). Hasta ahora, XNNPack obligaba a los desarrolladores a elegir entre la cuantización completa de enteros, donde tanto los pesos como las activaciones son int8, o el uso de coma flotante en precisión simple (fp32) o media (fp16). La nueva implementación de rango dinámico ofrece un punto intermedio que equilibra la velocidad de ejecución y la precisión del modelo.

Funcionamiento de la cuantización de rango dinámico

En el sistema de cuantización de rango dinámico, los pesos de los operadores Fully Connected y Convolution se convierten a enteros de 8 bits durante la fase de conversión del modelo. Sin embargo, el resto de los tensores se mantienen en formato float32. Durante la inferencia, las activaciones de las capas de coma flotante se transforman dinámicamente en enteros de 8 bits justo antes de pasar por los operadores optimizados. Este proceso calcula los parámetros de cuantización, específicamente la escala y el punto cero, basándose en el rango observado de las activaciones en tiempo real.

Esta metodología difiere de la cuantización completa en que los parámetros de escala no son fijos, sino que se adaptan dinámicamente. En los modelos totalmente cuantizados, estos valores se establecen durante la conversión utilizando un conjunto de datos representativo para calibrar el rango de las activaciones. Al calcularlos en tiempo de ejecución, la cuantización dinámica aprovecha mejor los 8 bits disponibles, lo que generalmente se traduce en una mayor precisión global del modelo en comparación con la cuantización estática, manteniendo una ganancia de rendimiento similar.

Otra diferencia técnica fundamental reside en la salida de los operadores. Mientras que en la cuantización completa el resultado es un entero de 8 bits, en la cuantización de rango dinámico la salida de los operadores Fully Connected y Convolution 2D se entrega en formato de coma flotante de 32 bits. Anteriormente, la inferencia de estos modelos se realizaba mediante los operadores nativos de TensorFlow Lite, pero ahora pueden aprovechar las implementaciones altamente optimizadas de XNNPack diseñadas específicamente para cada arquitectura de hardware.

Compatibilidad de hardware e implementación técnica

La compatibilidad con estas optimizaciones abarca una amplia gama de arquitecturas, incluyendo ARM, ARM64, x86 con instrucciones SSE, AVX y AVX512, además de WebAssembly. Google ha confirmado que estas mejoras son compatibles con los procesadores más recientes, como el Tensor G3 presente en el Pixel 8 o el Snapdragon 8 Gen 2 del OnePlus 11, basándose en la arquitectura ArmV9.

Para implementar esta funcionalidad, los desarrolladores deben convertir sus modelos de TensorFlow activando el flag de optimización converter.optimizations = [tf.lite.Optimize.DEFAULT]. Los modelos que ya hayan sido convertidos previamente utilizando cuantización de rango dinámico no requieren un nuevo proceso de conversión. Una ventaja competitiva de este método frente a la cuantización completa es que no requiere un conjunto de datos representativo y no falla si encuentra operadores no compatibles, lo que facilita su adopción por parte de desarrolladores que no son expertos en optimización de modelos.

Google ha anunciado que, a partir de la versión 2.17 de TensorFlow, la inferencia de XNNPack cuantizada dinámicamente estará activada por defecto en los binarios precompilados. Para quienes necesiten acceder a esta capacidad de forma inmediata, la funcionalidad ya se encuentra disponible en las compilaciones nocturnas (nightly builds) de TensorFlow.

Integración de precisión mixta con fp16

Además de la cuantización dinámica, XNNPack permite ahora combinar esta técnica con la inferencia de precisión mixta utilizando fp16 (coma flotante de media precisión). El formato fp16 utiliza la mitad de bits que el fp32 para almacenar un valor, reduciendo la precisión debido a una mantisa más corta (10 bits frente a 23 bits). No obstante, en dispositivos que cuentan con soporte hardware para fp16, el coste computacional de los operadores de coma flotante vectorizados se reduce a la mitad, ya que la CPU puede procesar el doble de datos por instrucción.

En este esquema de precisión mixta, los operadores Fully Connected y Convolution 2D pueden generar salidas en fp16 en lugar de fp32. Esto es especialmente beneficioso para modelos que incluyen operadores de coma flotante intensivos en cómputo, como Softmax o Batch Matrix Multiply, que pueden coexistir con la cuantización dinámica de los pesos para maximizar el rendimiento en el dispositivo.

Pruebas de rendimiento y resultados

Para validar estas mejoras, se han realizado pruebas de rendimiento con cuatro modelos públicos de visión computacional y generación de imágenes: EfficientNetV2 para clasificación de imágenes, Inception-v3 también para clasificación, Deeplab-v3 para segmentación semántica y Stable Diffusion para la generación de imágenes. En el caso de Stable Diffusion, la cuantización completa de enteros no era viable debido a la existencia de operadores no compatibles, lo que resalta la utilidad de la cuantización de rango dinámico como alternativa viable para modelos complejos.

Los resultados comparativos muestran que la implementación de XNNPack con cuantización de rango dinámico (QD8-F32 y QD8-F16) ofrece una aceleración considerable respecto al modelo base en float32. Específicamente, Google indica que el rendimiento de la inferencia en el backend de XNNPack se ha cuadruplicado en comparación con la línea base de precisión simple al añadir la compatibilidad para cuantización de rango dinámico en los operadores clave.

Impacto en el despliegue de IA on-device

Esta actualización tiene consecuencias directas en la capacidad de despliegue de aplicaciones de IA. Al reducir la dependencia de hardware de gama alta y eliminar la necesidad de procesos de calibración complejos, es posible llevar funciones avanzadas de inteligencia artificial a dispositivos más antiguos o con recursos limitados. La capacidad de ejecutar modelos con una latencia menor sin sacrificar drásticamente la precisión permite que la IA on-device sea más fluida y accesible.

En resumen, la integración de la cuantización de rango dinámico en XNNPack cierra la brecha entre la facilidad de uso de los modelos de coma flotante y la eficiencia de los modelos totalmente cuantizados. Al optimizar la ejecución en CPU y permitir la interoperabilidad con fp16, Google proporciona una herramienta más flexible para el despliegue de modelos en el ecosistema de TensorFlow Lite, optimizando la experiencia del usuario final mediante una ejecución más rápida y eficiente del procesamiento de datos en el dispositivo.