Resumen Técnico: FAR (Reemplazo de Atención Preservando Funciones para Inferencia Compatible con IMC)

1. Declaración del Problema

Las arquitecturas Transformer, particularmente aquellas que dependen de la autoatención, enfrentan cuellos de botella significativos en eficiencia al desplegarse en hardware emergente de Computación en Memoria (IMC), como las matrices cruzadas ReRAM. Si bien la autoatención destaca en procesadores digitales paralelos (GPUs/TPUs) al permitir la integración global de tokens, sus características computacionales entran en conflicto fundamental con el modelo de ejecución IMC:

• Multiplicaciones de Activación a Activación: La autoatención depende en gran medida de operaciones como $QK^\top$ y la normalización softmax, que requieren una mezcla dinámica y no local de activaciones, en lugar de las multiplicaciones vector-matriz (VMM) con pesos estacionarios que optimiza el hardware IMC.
• Sobrecarga de Memoria: La complejidad cuadrática ($O(T^2)$) de las interacciones entre tokens y la necesidad de leer repetidamente activaciones distribuidas espacialmente conducen a operaciones analógicas fragmentadas, conversiones extensas de Analógico a Digital/Digital a Analógico (ADC/DAC) y un pobre aprovechamiento de las matrices cruzadas.
• Limitaciones Actuales: Los aceleradores IMC existentes a menudo descargan la atención a unidades digitales, dejando el cuello de botella cuadrático sin resolver, o intentan mapear la atención directamente a matrices analógicas sin abordar los costos inherentes de movimiento de datos.

Simultáneamente, análisis recientes sugieren que las capas de atención exhiben una redundancia sustancial durante la inferencia, con mapas de atención volviéndose más suaves y comprimibles a través de la profundidad. Esto motiva la hipótesis de que el rol funcional de la atención puede aproximarse mediante módulos secuenciales que se alinean mejor con la localidad y la estacionariedad de los pesos en IMC.

2. Metodología: El Marco FAR

Los autores proponen FAR (Reemplazo de Atención Preservando Funciones), un marco que reemplaza todos los bloques de autoatención en un Transformer de Visión preentrenado (específicamente la familia DeiT) con módulos secuenciales, concretamente LSTMs Bidireccionales de Múltiples Cabezas (BiLSTMs). El enfoque está diseñado como una conversión arquitectónica "lista para usar" que no requiere entrenamiento desde cero.

2.1 Estrategia de Reemplazo

• Sustitución por Bloques: Cada bloque de atención $A_l$ en un Transformer preentrenado se reemplaza por un sustituto aprendible $A'_l$ (un BiLSTM de múltiples cabezas). Todos los demás componentes (bloques MLP, incrustaciones de parches, normalización de capas) permanecen congelados.
• Equivalencia Funcional mediante Distilación: El reemplazo se supervisa mediante distilación por capas. Para cada capa $l$, una pérdida de similitud ($L_{sim}$) minimiza la diferencia entre la salida de la cabeza de atención original y la nueva cabeza BiLSTM:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  Esto asegura que el sustituto replique la estructura representacional del profesor.
• Protocolo de Entrenamiento: El proceso involucra dos etapas:
  1. Fase de Distilación: Solo los módulos sustitutos BiLSTM se entrenan utilizando una combinación de pérdida de similitud y pérdida de tarea (clasificación), mientras el resto del modelo permanece congelado.
  2. Fase de Ajuste Fino: Todo el modelo se ajusta finamente conjuntamente para recuperar cualquier caída residual de precisión.

2.2 Arquitectura BiLSTM de Múltiples Cabezas

Para preservar el sesgo inductivo de la atención de múltiples cabezas, el sustituto se estructura como una BiLSTM de Múltiples Cabezas:

• Proyección y División: La incrustación de entrada se proyecta y divide en $N$ subespacios (coincidiendo con el número de cabezas de atención del profesor).
• Procesamiento Bidireccional: Cada subespacio se procesa mediante una cabeza BiLSTM separada, capturando dependencias de tokens tanto hacia adelante como hacia atrás.
• Concatenación y Proyección: Las salidas se concatenan y se proyectan de nuevo a la dimensión original.
• Compatibilidad con IMC: A diferencia de la atención, esta estructura depende de la reutilización recurrente de pesos y del flujo de datos secuencial, eliminando las interacciones de activación entre todos los pares. Esto se mapea naturalmente a matrices cruzadas ReRAM donde los pesos son estacionarios y los datos fluyen secuencialmente.

2.3 Poda Estructurada para Restricciones IMC

Para adaptarse aún más a matrices IMC con recursos limitados, los autores aplican poda estructurada a las dimensiones ocultas de las cabezas BiLSTM:

• Extensión de DeepHoyer: Extienden el marco DeepHoyer para apuntar a unidades ocultas de LSTM. Dado que las puertas de LSTM (entrada, olvido, celda, salida) comparten unidades ocultas, podar una sola unidad requiere una eliminación coordinada a través de todas las matrices de peso asociadas ($W_{ih}$, $W_{hh}$ y pesos de proyección).
• Esparsidad de Grupo: Se aplica una medida de esparsidad Hoyer diferenciable a grupos de pesos correspondientes a cada unidad oculta, fomentando que la red aprenda qué unidades son redundantes.
• Compresión Adaptativa por Capa: La poda se realiza independientemente por cabeza y dirección, permitiendo que el modelo adapte su capacidad a la redundancia específica de cada capa, resultando en módulos compactos que se ajustan a los tamaños de las matrices IMC.

3. Resultados Clave

El marco se evaluó en la familia DeiT (Tiny, Small, Base) en ImageNet y varias tareas posteriores (CIFAR-10/100, Coches, Flores, iNaturalist).

• Preservación de Precisión:
  • DeiT-Tiny: FAR-Tiny en realidad superó al DeiT-Tiny de referencia en un +1.2% en ImageNet (73.3% vs 72.2%), sugiriendo que el sustituto secuencial actúa como un sesgo inductivo beneficioso para modelos más pequeños.
  • Modelos Más Grandes: Para DeiT-Small y Base, FAR mantuvo una precisión competitiva con una brecha modesta de aproximadamente 1–2% en comparación con el profesor, demostrando que el rol funcional de la atención es en gran medida reproducible mediante recurrencia secuencial.
  • Transferencia a Tareas Posteriores: Los modelos FAR mantuvieron una fuerte transferibilidad a tareas de clasificación posteriores, indicando que los mapeos de secuencia a secuencia aprendidos se generalizan efectivamente.
• Robustez ante la Poda:
  • La aplicación de poda estructurada redujo la dimensión oculta promedio en un 40–60% con una caída adicional de precisión de solo 1–2%.
  • Los patrones de poda fueron no uniformes, con mayor retención en las capas intermedias y una poda más agresiva en los extremos de la red, reflejando la heterogeneidad interna en la redundancia representacional.
• Especialización por Cabeza: La visualización de los mapas de interacción de tokens mostró que las cabezas FAR heredaron con éxito los roles semánticos distintos de las cabezas de atención originales, concentrándose en regiones específicas de objetos en lugar de colapsar en patrones redundantes.

Análisis de Eficiencia de Hardware

Utilizando un modelo analítico de IMC ReRAM, los autores estimaron la latencia y la energía:

• Latencia: Se estima que FAR en IMC ReRAM es ~18 veces más rápido que DeiT en IMC ReRAM (debido a evitar la sobrecarga cuadrática de la atención) y ~400 veces más rápido que DeiT en una GPU.
• Energía: Se estima que FAR en IMC ReRAM es ~3 veces más eficiente energéticamente que DeiT en IMC ReRAM y ~150 veces más eficiente que DeiT en una GPU.
• Mecanismo: Estas ganancias provienen de eliminar las multiplicaciones de activación a activación, reducir las conversiones ADC/DAC y permitir la ejecución con pesos estacionarios. Los autores señalan que excluir los costos de softmax (que son cuellos de botella significativos para la atención en IMC) hace que estas estimaciones sean conservadoras para FAR.

4. Significado y Afirmaciones

El artículo afirma que FAR proporciona una alternativa práctica y eficiente para el hardware para desplegar Transformers preentrenados en aceleradores de borde basados en IMC sin necesidad de entrenar modelos desde cero.

• Viabilidad del Reemplazo Total: El trabajo demuestra que la atención puede reemplazarse de extremo a extremo en tiempo de inferencia manteniendo una precisión práctica, desafiando la noción de que las interacciones de tokens entre todos los pares son estrictamente necesarias para tareas de visión de alto rendimiento.
• Alineación con Hardware: Al sustituir la atención con módulos secuenciales y con pesos estacionarios, FAR alinea el flujo de datos computacional con las fortalezas físicas de las matrices cruzadas ReRAM, reduciendo significativamente el tráfico de memoria y la latencia.
• Preservación del Conocimiento Preentrenado: El marco demuestra que las ricas representaciones semánticas aprendidas durante el preentrenamiento a gran escala pueden preservarse y transferirse a arquitecturas estructuralmente diferentes y compatibles con IMC mediante distilación por bloques.

En resumen, FAR ofrece una vía para reestructurar la inferencia de Transformers en torno a módulos secuenciales compatibles con IMC, cerrando la brecha entre el dominio de los modelos basados en atención y las restricciones del hardware emergente centrado en memoria.