ScalePredictor: Instance-aware Scale Learning for Accurate Quantization of Vision Transformers

Resumen Técnico: ScalePredictor

Planteamiento del Problema

Los Vision Transformers (ViTs) han alcanzado un rendimiento de vanguardia en la visión por computadora, pero enfrentan desafíos significativos de despliegue en dispositivos periféricos (edge) con recursos limitados debido a su intensidad computacional. La Cuantización Post-Entrenamiento (PTQ) ofrece una solución al comprimir los modelos a formatos de pocos bits (por ejemplo, 8 bits o menos) sin necesidad de reentrenamiento. Sin embargo, los métodos de PTQ existentes dependen predominantemente de la cuantización estática, donde un único conjunto de parámetros de cuantización (escalas) se aplica uniformemente a todas las instancias de entrada.

Este paradigma estático es inherentemente subóptimo para los ViTs porque las imágenes naturales exhiben una diversidad sustancial, lo que conduce a variaciones significativas en las distribuciones de activación entre diferentes muestras. Esta heterogeneidad es particularmente perjudicial en configuraciones de ultra bajo bit (por ejemplo, W3/A3), donde las escalas fijas no logran capturar la dinámica específica de cada instancia. Por el contrario, las estrategias de cuantización dinámica que adaptan los parámetros por muestra suelen requerir reducciones de tensores completos en tiempo de ejecución (cálculos de min/max) para cada capa. Estas operaciones introducen barreras de sincronización severas y sobrecargas limitadas por la memoria, lo que socava drásticamente la velocidad de inferencia y anula los beneficios de eficiencia de la cuantización.

Metodología: ScalePredictor

Los autores proponen ScalePredictor, un marco de cuantización dinámica que logra un aprendizaje de escala consciente de la instancia con una sobrecarga computacional mínima. La idea central que impulsa el método es una correlación descubierta: las escalas de cuantización óptimas para las capas profundas dependen fuertmente de la magnitud de la activación observada en las capas superficiales.

El marco consta de tres componentes principales:

1. Extracción de Rango Robusto:
   En lugar de calcular valores min/max globales (sensibles a los valores atípicos) o realizar un costoso ordenamiento de percentiles, ScalePredictor extrae un "rango robusto" ($r(x)$) de la salida del Patch Embedding (la capa más superficial).

   • El tensor aplanado se particiona en $K$ fragmentos (chunks) no superpuestos.
   • Se calculan los valores extremos locales (min/max) para cada fragmento.
   • Estos extremos locales se promedian para derivar un rango robusto global.
   • Este enfoque actúa como una aproximación de percentil de grano grueso, suprimiendo eficazmente la influencia de los valores atípicos mientras mantiene una complejidad de $O(N \cdot C)$ y evita el ordenamiento.

2. Predictor Polinomial Motivado por Taylor:
   Un módulo polinomial aprendible mapea el rango robusto extraído $r(x)$ a las escalas de cuantización óptimas para todas las capas subsiguientes.

   • Basándose en la observación de que la relación entre las estadísticas superficiales y las escalas óptimas profundas sigue una tendencia suave y no lineal, el método modela la escala $s_l(x)$ para la capa $l$ como una función polinómica de $r(x)$:
     $$s_l(x) = \sum_{k=0}^{d} c_{l,k} \cdot (r(x))^k$$
   • Los coeficientes $c_{l,k}$ se aprenden durante una fase de calibración utilizando una pérdida de reconstrucción por bloques.
   • Esto permite predecir escalas específicas de la instancia para todas las capas en una sola pasada, sin acceder a los tensores de activación profundos.

3. Calibración e Inferencia:

   • Calibración: Los coeficientes polinómicos se optimizan utilizando un pequeño conjunto de calibración. El término constante ($c_{l,0}$) se inicializa con la escala estática óptima para asegurar la estabilidad, mientras que los términos de orden superior se aprenden para capturar las variaciones de la instancia.
   • Inferencia: Para una nueva entrada, el rango robusto se calcula una vez en la capa de Patch Embedding. Todas las escalas de las capas subsiguientes se generan mediante la evaluación polinómica ($O(1)$ FLOPs), eliminando la necesidad de reducciones por capa en tiempo de ejecución.

Principales Contribuciones

1. PTQ Consciente de la Instancia para ViTs: ScalePredictor se presenta como el primer método de cuantización post-entrenamiento para ViTs que aprovecha la cuantización dinámica para abordar la heterogeneidad de la distribución a nivel de instancia sin reentrenamiento.
2. Mecanismo de Aprendizaje de Escala Eficiente: La integración de un mecanismo de extracción de rango robusto (promedio de min/max basado en fragmentos) con una proyección polinomial inspirada en Taylor permite la predicción de escalas precisas y en tiempo real con una sobrecarga de inferencia insignificante.
3. Descubrimiento de la Correlación Superficial-Profunda: El trabajo revela una dependencia funcional consistente donde las estadísticas de activación de las capas superficiales (específicamente de los Patch Embeddings) pueden predecir eficazmente las escalas de cuantización óptimas para los bloques transformadores más profundos en diversas arquitecturas (ViT, DeiT, Swin).

Resultados Experimentales

Se realizaron experimentos extensos en el conjunto de datos ImageNet a través de siete arquitecturas de ViT (ViT-S/B, DeiT-T/S/B, Swin-S/B) y cuatro líneas base de PTQ (BRECQ, QDrop, AdaLog, I&S-ViT).

• Mejoras en la Precisión: ScalePredictor supera consistentemente a las líneas base de PTQ estática, particularmente en regímenes de pocos bits.
  • Las ganancias notables incluyen una mejora del 10.58% (de 33.57% a 44.15%) en Swin-Small usando I&S-ViT en W2/A3.
  • También se observaron mejoras significativas en W3/A3, como una ganancia del 8.71% en ViT-Small usando BRECQ.
  • Incluso en casos donde las líneas base estáticas colapsaron (por ejemplo, BRECQ en ViT-Small en W2/A3 con 1.10% de precisión), ScalePredictor recuperó el rendimiento a 4.15%.
• Comparación con Cuantización Dinámica JIT: A diferencia de los métodos de cuantización dinámica Just-In-Time (JIT) (por ejemplo, JIT Min/Max), que degradan severamente la precisión al romper la relación calibrada entre peso y activación, ScalePredictor preserva la consistencia de la calibración permitiendo la adaptabilidad.
• Eficiencia:
  • Latencia: En una NVIDIA RTX 3090, ScalePredictor introduce solo un 0.6% de sobrecarga de latencia de extremo a extremo en comparación con el PTQ estático.
  • Cálculo de Escalas: El método es 40–52 veces más rápido que JIT Min/Max y 507–1025 veces más rápido que JIT Percentile para calcular las escalas de cuantización, ya que evita las reducciones de tensores por capa.

Significado y Reivindicaciones

El artículo afirma que ScalePredictor logra cerrar la breancia entre la precisión de la cuantización dinámica y la eficiencia de la cuantización estática. Al aprovechar la correlación oculta entre las distribuciones de las capas superficiales y las escalas óptimas de las capas profundas, el método evita la costosa calibración "just-in-time" requerida por los enfoques dinámicos tradicionales.

Los autores enfatizan que ScalePredictor ofrece un equilibrio favorable entre precisión y eficiencia, lo que lo convierte en una solución viable para el despliegue de Vision Transformers de bajo bit en dispositivos periféricos. El método está diseñado para integrarse sin problemas en los marcos de trabajo de PTQ existentes, requiriendo solo la adición de un módulo predictor ligero y un paso de extracción de rango robusto, sin alterar la arquitectura subyacente de la red ni requerir un reentrenamiento completo.