Resumen Técnico: LL-ViT – Vision Transformers Desplegables en el Borde con Neuronas de Tabla de Búsqueda (Look Up Table)

Declaración del Problema

Los Vision Transformers (ViTs) han alcanzado un rendimiento de vanguardia en visión por computadora, pero enfrentan barreras significativas para su despliegue en dispositivos de borde (edge), particularmente en matrices de puertas programables en campo (FPGAs). Su gran huella computacional, de memoria y de energía excede las capacidades de las plataformas de borde comerciales. Por ejemplo, variantes compactas como DeiT-T requieren aproximadamente 20 MB de almacenamiento de pesos, superando la memoria interna de muchas FPGAs.

Las estrategias de optimización existentes, como la carga de pesos fuera del chip y la cuantización agresiva, introducen compensaciones: el acceso frecuente fuera del chip aumenta el consumo de energía, mientras que la cuantización agresiva suele degradar la precisión del modelo. Además, soluciones híbridas recientes que utilizan ejecución en pipeline en plataformas de alto nivel (por ejemplo, AMD/Xilinx Versal) logran un alto rendimiento, pero a niveles de potencia (por ejemplo, ~40W) no aptos para escenarios de borde alimentados por batería.

Trabajos recientes sobre redes basadas en Lógica y Tablas de Búsqueda (LUT) (por ejemplo, LogicNets, NeuraLUT, Differentiable Weightless Neural Networks) ofrecen modelos con huellas de memoria y cómputo reducidas, pero históricamente se han limitado a conjuntos de datos pequeños y tareas simples. Estos modelos suelen tener un desempeño deficiente en benchmarks estándar de visión por computadora como CIFAR-10/100 y no se han integrado con éxito en la arquitectura compleja de los Vision Transformers.

Metodología

Los autores proponen LL-ViT (Learned-LUT based Vision Transformers), un novedoso codiseño de algoritmo-hardware que integra neuronas basadas en LUT directamente en la arquitectura del transformer para optimizar la inferencia en el borde.

1. Análisis de Carga de Trabajo y Motivación

Los autores caracterizan arquitecturas típicas de ViT (específicamente DeiT-T) e identifican que las capas de Perceptrón Multicapa (MLP) (mezcladores de canales) son el cuello de botella principal. Estas capas contribuyen con más del 60% de los pesos totales del modelo y el 55% de las operaciones de Multiplicación-Acumulación (MAC) totales. Esta observación motiva el reemplazo del bloque MLP estándar por una alternativa más eficiente, manteniendo el mezclador de tokens de Auto-Atención de Múltiples Cabezales (MHA), que es menos dominante en cuanto a pesos.

2. Mezclador de Canales Basado en LUT Aprendible

En lugar de simplemente reemplazar multiplicaciones por búsquedas en tablas post-entrenamiento, los autores diseñan un mezclador de canales basado en LUT aprendible que se entrena de extremo a extremo.

• Arquitectura: El bloque propuesto reemplaza el MLP estándar de dos capas. Utiliza codificación de termómetro para convertir las activaciones de entrada en secuencias de bits, que luego son procesadas a través de capas configurables de neuronas LUT.
• Diferenciabilidad: Para permitir la retropropagación, los autores emplean técnicas de Diferencia Finita Extendida (de Differentiable Weightless Neural Networks) para aproximar los gradientes de las entradas de la LUT.
• Sumatoria Condicional: Una novedosa capa de sumatoria condicional sigue a las neuronas LUT. Esta añade los valores codificados a la salida solo si la salida de la neurona LUT correspondiente está activa (1). Esta capa permite que la red produzca las salidas de valor real requeridas por el transformer mientras mantiene una ruta de inferencia libre de multiplicaciones. Durante el entrenamiento, estos valores codificados se aprenden en precisión completa (fp32) y se cuantizan post-entrenamiento (por ejemplo, a int4) para coincidir con el resto de la red.

3. Diseño del Acelerador FPGA

Los autores codiseñan un acelerador de hardware específicamente para LL-ViT:

• Elementos de Procesamiento (PE): Se diseña un PE dedicado para el mezclador de canales basado en LUT. Este procesa las activaciones de entrada fila por fila de forma segmentada (pipelined) utilizando buffers ping-pong para evitar paradas.
• Mapeo de Recursos: Las neuronas LUT se mapean directamente a las celdas LUT dentro de la estructura lógica de la FPGA, eliminando la necesidad de DSPs o BRAMs para el almacenamiento de pesos en estas capas. Esto contrasta con los aceleradores convencionales que requieren DSPs para la multiplicación y BRAMs para el almacenamiento de pesos.
• Pesos Estacionarios: Todos los pesos del modelo se almacenan en la memoria interna en BRAMs dedicadas, eliminando el acceso a la memoria fuera del chip durante la inferencia.
• Mezclador de Tokens: El bloque MHA se implementa utilizando una matriz sistólica para multiplicaciones de matrices, interconectada con el PE del mezclador de canales basado en LUT.

Principales Contribuciones

1. Arquitectura Novedosa: Propuesta de un bloque mezclador de canales basado en LUT aprendible que se integra sin fisuras en los Vision Transformers, reemplazando las capas MLP computacionalmente pesadas.
2. Modelo Optimizado para el Borde: Desarrollo de un ViT optimizado para el borde (LL-ViT) que combina la mezcla de tokens por auto-atención con la mezcla de canales libre de multiplicaciones, reduciendo significativamente la demanda computacional y el consumo de memoria.
3. Codiseño de Hardware: Diseño de un acelerador FPGA de bajo consumo que evita el movimiento de pesos fuera del chip y aprovecha los recursos nativos de LUT de las FPGAs para el mezclador de canales.
4. Desempeño en Tareas Complejas: Extensión de las redes neuronales basadas en LUT desde conjuntos de datos de pequeña escala hacia tareas complejas de visión por computadora (CIFAR-10/100, TinyImageNet), demostrando que las neuronas LUT pueden aprender invariancia posicional e interacciones de características complejas.

Resultados Experimentales

Los autores evaluaron LL-ViT en una FPGA AMD Xilinx Virtex (xcvu9p) contra una línea base totalmente cuantizada (INT8) (I-ViT-T) y otros trabajos previos.

• Precisión: LL-ViT logra precisiones comparables a la línea base:
  • CIFAR-10: 95.5% (vs. 95.4% de la línea base)
  • CIFAR-100: 78.8% (vs. 79.2% de la línea base)
  • TinyImageNet: 60.9% (vs. 60.4% de la línea base)
  • Flowers-102: 91.6% (vs. 91.3% de la línea base)
• Tamaño del Modelo y Cómputo:
  • Reducción del 60% en los pesos del modelo (1.93 MB vs. 5.06 MB).
  • Reducción de >50% en las operaciones MAC (0.65 GMACs vs. 1.31 GMACs).
  • El tamaño reducido permite que todo el modelo quepa dentro de la memoria interna.
• Eficiencia:
  • Mejora de 1.9× en la eficiencia energética (2.14 mJ/inferencia vs. 4.05 mJ para la línea base).
  • Latencia 1.3× menor (5.33 ms vs. 6.93 ms).
  • Rendimiento (Throughput): Alcanza 1083 FPS con un presupuesto de potencia de 10.9W.
• Comparación: LL-ViT supera a trabajos previos basados en LUT (por ejemplo, TreeLUT, DiffLogicNet) que reportan precisiones significativamente menores (40–80% en CIFAR-10) y demuestra una eficiencia superior en comparación con otros aceleradores de ViT (por ejemplo, HG-PIPE, ViTA) con una fracción de la potencia.

Significado y Reivindicaciones

El artículo afirma que LL-ViT representa un paso significativo para hacer que los Vision Transformers sean viables para el despliegue en el borde. Al dirigirse a las capas MLP —la principal fuente de sobrecarga de memoria y cómputo— con neuronas LUT aprendibles, los autores demuestran que es posible:

1. Reducir drásticamente el tamaño del modelo y el consumo de energía sin sacrificar la precisión en benchmarks de visión estándar.
2. Utilizar plenamente los recursos de la FPGA al mapear los cálculos neuronales directamente a las celdas LUT, eliminando la necesidad de acceso a la memoria fuera del chip y reduciendo la dependencia de los costosos DSPs.
3. Cerrar la brecha entre las redes neuronales basadas en LUT eficientes y las arquitecturas complejas de transformers, probando que las neuronas LUT pueden entrenarse para manejar las complejidades de las tareas de visión anteriormente reservadas para redes densas y pesadas en multiplicaciones.

Los autores posicionan a LL-ViT como una "alternativa ligera prometedora" a los ViTs tradicionales, allanando el camino para "tiny ViTs" que puedan ser desplegados en dispositivos de borde con recursos limitados y alimentados por baterías. Observan que, aunque la técnica se demostró en una línea base de tiny ViT, se espera que los beneficios estructurales escalen a variantes más grandes.