Resumen Técnico: FTerViT: Vision Transformer Completamente Ternario

Enunciado del Problema

Los Vision Transformers (ViT) son clasificadores de imágenes potentes, pero generalmente no son adecuados para dispositivos de clase microcontrolador debido a su considerable huella de memoria. Un ViT compacto estándar como DeiT-Small requiere aproximadamente 88.3 MB en precisión FP32, mientras que las unidades de microcontrolador (MCU) típicas ofrecen solo unos pocos megabytes de RAM externa. Esta discrepancia impide el despliegue de aplicaciones de visión siempre activas y de bajo consumo en dispositivos periféricos sin descarga a la nube.

Aunque la cuantización ternaria (restringir los pesos a $\{-1, 0, +1\}$) ofrece una compresión teórica de 16$\times$ sobre FP32, los métodos existentes del estado del arte solo ternarizan las capas del codificador del transformador. Mantienen la incrustación de parches, los parámetros de LayerNorm y las cabezas de clasificador en precisión completa (FP32) o enteros de menor bit (INT8). En modelos compactos, estos componentes restantes de precisión completa se convierten en el cuello de botella dominante, consumiendo hasta el 38% del tamaño del modelo a pesar de representar una pequeña fracción de los parámetros. Además, los intentos previos de ternarizar completamente estos componentes sensibles (por ejemplo, TerViT) resultaron en caídas significativas de precisión (de hasta 22.4 puntos porcentuales), en gran parte debido a la conocida sensibilidad de las primeras y últimas capas a la cuantización.

Metodología

Los autores proponen FTerViT, el primer Vision Transformer completamente ternarizado donde todas las matrices de pesos y los parámetros de normalización están restringidos a $\{-1, 0, +1\}$. El enfoque consta de dos pilares principales: operadores ternarios novedosos y un protocolo de entrenamiento especializado.

1. Primitivas Ternarias

Para eliminar los cuellos de botella FP32, los autores reemplazan los componentes estándar con módulos ternarios dedicados:

• TernaryBitLinear: Utilizado para capas totalmente conectadas. Cuantiza los pesos a $\{-1, 0, +1\}$ utilizando escalado de media absoluta por tensor ($absmean$) y redondeo acotado. Las activaciones se cuantizan a enteros de 8 bits utilizando escalado de máximo absoluto por token.
• TernaryBitConv2d: Diseñado específicamente para la capa de incrustación de parches. Aplica escalado por canal para manejar magnitudes de filtros heterogéneas, abordando la sensibilidad de la capa de incrustación de parches, que trabajos previos identificaron como el componente más crítico.
• TernaryLayerNorm: Ternariza los parámetros afines aprendibles ($\gamma$ y $\beta$) de LayerNorm utilizando el mismo esquema $absmean$. Las estadísticas por token (media y varianza) y los sesgos permanecen en FP32 ya que su conteo de parámetros es insignificante, pero su ternarización destruiría la información posicional.

2. Protocolo de Entrenamiento

El modelo se entrena utilizando una estrategia de Distilación de Conocimiento (KD) en dos fases:

• Fase 1 (Entrenamiento): El estudiante ternario se inicializa a partir de un maestro FP32 preentrenado y se entrena utilizando Distilación Consciente de Cuantización (QAD) con una tasa de aprendizaje de $10^{-4}$ (decaimiento coseno) durante 250 épocas. La función de pérdida utiliza únicamente el término de divergencia KL hacia adelante entre las logit del estudiante y del maestro, evitando la entropía cruzada que puede entrar en conflicto con KL en redes de bajo bit.
• Fase 2 (Recuperación): Una fase de ajuste fino ligera donde la tasa de aprendizaje se reinicia en $10^{-5}$ (decaimiento coseno) durante 10 épocas. Esta fase permite la recalibración de escalas y la recuperación de la precisión perdida durante la saturación inicial de la Fase 1.

Contribuciones Clave

1. Primer ViT Completamente Ternario: Los autores demuestran que los componentes más frágiles de los ViT (incrustación de parches, LayerNorms y cabeza de clasificador) pueden ternarizarse exitosamente a $\{-1, 0, +1\}$ cuando se entrenan con distilación de conocimiento, empujando los límites de lo que anteriormente se consideraba cuantizable.
2. Operadores Novedosos: Introducción de TernaryBitConv2d con escalado por canal y TernaryLayerNorm para manejar la sensibilidad específica de las capas no codificadoras.
3. Estrategia de Entrenamiento: Una estrategia de entrenamiento simple pero efectiva en dos fases (entrenamiento con alta tasa de aprendizaje seguido de recuperación con baja tasa de aprendizaje) que permite una convergencia más rápida y una precisión final significativamente mayor en comparación con enfoques previos de ViT ternario.
4. Implementación en Dispositivo: La primera implementación C independiente de un Vision Transformer ternario en un microcontrolador ESP32-S3 de doble núcleo con restricciones de recursos (8 MB de PSRAM).

Resultados Experimentales

El método se evaluó en ImageNet-1K y otras pruebas de referencia:

• Precisión y Compresión:

  • FTerViT-DeiT-III-S (384$\times$384): Logra una precisión top-1 del 82.43% a 6.09 MB (compresión de 14.6$\times$), perdiendo solo 2.42 puntos porcentuales en comparación con su maestro FP32. Esto supera a los métodos previos de ViT ternario (por ejemplo, TerViT) en hasta 8 puntos porcentuales.
  • FTerViT-DeiT-III-S (224$\times$224): Logra un 79.64% top-1 a 5.81 MB.
  • FTerViT-DeiT-S (224$\times$224): Logra un 77.47% top-1 a 5.81 MB, superando a TerViT-DeiT-S en 3.27 puntos porcentuales.
  • Pruebas de referencia CIFAR: Las variantes DeiT-Tiny logran un 97.43% (CIFAR-10) y un 86.01% (CIFAR-100) con solo 1.53 MB de almacenamiento.

• Análisis de Sensibilidad de Componentes:

  • Los análisis de primer orden de Taylor y de traza de Hessian confirman que los LayerNorms y las incrustaciones de parches dominan la importancia del modelo (34–39% y 4.6% respectivamente) a pesar de ocupar <0.2% de los parámetros. Esto explica por qué los métodos previos los dejaron en mayor precisión y por qué su ternarización exitosa es crítica.

• Despliegue en Dispositivo (ESP32-S3):

  • El modelo de 5.81 MB cabe completamente en los 8 MB de PSRAM del ESP32-S3, dejando ~4.5 MB libres para búferes de cámara y LCD.
  • Rendimiento: Un paso hacia adelante tarda 21.06 segundos.
  • Eficiencia: En comparación con FP32, la inferencia ternaria empaquetada de 2 bits es 1.71$\times$ más rápida (804 ms vs. 1376 ms para una capa específica) y reduce el consumo de energía por encima del estado inactivo en un 55% (59 mJ vs. 130 mJ).

Significado

El artículo afirma que FTerViT reduce sustancialmente la brecha entre la precisión estándar de ViT y las estrictas restricciones de memoria del despliegue a escala de microcontrolador. Al ternarizar con éxito toda la arquitectura, incluidas las incrustaciones de parches y las capas de normalización previamente "intocables", los autores permiten el despliegue de modelos de visión de alta precisión en MCUs comerciales y de bajo costo ($10). El trabajo valida que con la estrategia de entrenamiento correcta (KD de misma arquitectura y recuperación en dos fases), la sensibilidad de estos componentes puede gestionarse, haciendo que los ViT completamente ternarios sean una realidad práctica para la IA en el borde. Los autores señalan que la brecha de precisión con FP32 escala inversamente con la capacidad del modelo, donde los modelos más grandes (DeiT-III-S384) sufren menos pérdida de precisión relativa que los más pequeños.