Beyond Self-Attention: Sub-Quadratic Vision Transformers for Fast Image Captioning

Resumen Técnico: Más allá de la autoatención: Vision Transformers subcuadráticos para el subtitulado rápido de imágenes

Planteamiento del problema
El subtitulado de imágenes requiere la generación de descripciones textuales coherentes y semánticamente significativas a partir de contenido visual, tendiendo un puente entre la visión por computadora y el procesamiento del lenguaje natural. Aunque las arquitecturas basadas en Transformers han avanzado significativamente en el campo al capturar dependencias contextuales de largo alcance mediante la autoatención, enfrentan dos limitaciones principales. Primero, los mecanismos de autoatención estándar computan interacciones por pares entre todos los parches de la imagen, lo que resulta en una complejidad computacional cuadrática ($O(n^2)$). Esto crea una sobrecarga sustancial de memoria y computación, particularmente para imágenes de alta resolución y aplicaciones a gran escala. Segundo, los enfoques existentes suelen sufrir de una falta de representaciones de características locales ricas y de redundancia computacional, ya que procesan todos los parches por igual, incluso cuando muchos carecen de información semánticamente crítica.

Metodología
La arquitectura propuesta aborda estas limitaciones reestructurando el marco de trabajo del Vision Transformer (ViT) para reemplazar el mecanismo convencional de autoatención con un enfoque de agrupamiento probabilístico basado en un Modelo de Mezcla Gaussiana (GMM). El sistema opera en dos etapas principales:

1. Codificación de imagen (GMM-Transformer):

   • Incrustación de parches (Patch Embedding): Las imágenes de entrada se dividen en parches no superpuestos, se proyectan en incrustaciones de $D$ dimensiones y se aumentan con codificaciones posicionales.
   • Agrupamiento GMM ponderado por contexto: En lugar de computar puntuaciones de atención por pares, el modelo trata las incrustaciones de los parches como variables aleatorias continuas generadas por una mezcla de $K$ distribuciones gaussianas multivariantes ($K \ll N$).
   • Pesos de contexto aprendibles: Una subcapa asigna pesos espaciales aprendibles a los parches, modulados por propiedades estadísticas (media y varianza), lo que permite al modelo priorizar regiones semánticamente significativas.
   • Algoritmo EM: Se emplea el algoritmo de Esperanza-Maximización (EM) para realizar un agrupamiento suave (soft-clustering). En el paso E, el modelo computa las "responsabilidades" (probabilidades de asignación suave) que indican el grado en que cada parche pertenece a un grupo. En el paso M, se actualizan los parámetros de los grupos (medias, covarianzas y coeficientes de mezcla).
   • Reducción de complejidad: Este proceso reemplaza la matriz de atención $N \times N$ con una matriz de responsabilidad $N \times K$, reduciendo la complejidad computacional de cuadrática $O(n^2)$ a casi lineal $O(nK)$.
   • Agregación: Las características refinadas de los parches se procesan a través de redes de alimentación hacia adelante (feed-forward) con conexiones residuales y normalización de capa. Finalmente, una operación de Agregación Global Promedio (Global Average Pooling) agrega los centroides de los grupos en un vector de características globales compacto.

2. Decodificación de lenguaje (Decodificador unificado basado en GPT):

   • El decodificador utiliza un Transformer de estilo GPT autorregresivo.
   • Atención unificada: A diferencia de los modelos encoder-decoder tradicionales que separan la autoatención enmascarada y la atención cruzada, esta arquitectura las integra en un único bloque de atención multi-cabezal.
   • Mecanismo: Los tokens textuales sirven como vectores de consulta (query), mientras que la secuencia combinada de tokens visuales (centroides de los grupos) y los tokens textuales previos sirven como vectores de clave (key) y valor (value).
   • Enmascaramiento causal: Una máscara causal especializada asegura que los tokens de texto atiendan a todos los tokens visuales pero solo a los tokens de texto previamente generados, manteniendo las propiedades de generación autorregresiva.
   • Generación: La salida se proyecta a través de una capa Softmax para predecir el siguiente token en la secuencia de forma iterativa.

Contribuciones clave

• Reducción de complejidad: La arquitectura reduce la complejidad temporal del codificador del transformer de cuadrática a casi lineal ($O(nK)$) al sustituir la autoatención por un agrupamiento suave basado en GMM, manteniendo la capacidad de generar subtítulos contextualmente coherentes.
• Diseño de decodificador eficiente: Se emplea un decodificador personalizado basado en GPT que unifica la autoatención enmascarada y la atención cruzada en un solo bloque computacional, mejorando la interacción multimodal y reduciendo la redundancia.
• Validación empírica: Una evaluación exhaustiva en el conjunto de datos Flickr30k demuestra el rendimiento competitivo del modelo frente a trabajos actuales del estado del arte (SOTA), equilibrando la eficiencia computacional con la calidad del subtitulado.

Resultados experimentales
El modelo fue evaluado en el conjunto de datos Flickr30k (usando la división de Karpathy) con una configuración de 8 capas de transformer, 4 cabezales de atención y una dimensión oculta de 768.

• Métricas de rendimiento: El modelo propuesto alcanzó puntuaciones máximas en métricas semánticas avanzadas: BLEU-4 (0.364), METEOR (0.368) y SPICE (0.199). Estos resultados indican una estructura de frase y un mapeo de grafos de escena superiores en comparación con varios modelos SOTA (por ejemplo, TAKG, GNN-Visual, Scene Graph).
• Comparación con líneas base: Aunque el modelo queda por detrás del modelo ADA en métricas de precisión de línea base (BLEU-1: 0.733 vs. 0.794; ROUGE: 0.568 vs. 0.608; CIDEr: 0.723 vs. 0.946), supera a otras líneas base como TopDown y GNN-Visual. Los autores señalan que la arquitectura favorece el contexto profundo sobre la simple replicación.
• Eficiencia: El modelo opera con un costo computacional de aproximadamente 21.17 GFLOPs y contiene alrededor de 32.8 millones de parámetros entrenables. El análisis de la matriz de confusión mostró altas tasas de verdaderos positivos (45,580) y bajos recuentos de falsos positivos/negativos, lo que indica una fuerte capacidad de clasificación.

Significancia y afirmaciones
El artículo afirma que el marco propuesto aborda con éxito la ineficiencia computacional inherente a los Vision Transformers estándar. Al aprovechar el Modelo de Mezcla Gaussiana y el algoritmo de Esperanza-Maximización, el modelo logra una reducción significativa de la complejidad computacional de cuadrática a casi lineal sin sacrificar la capacidad de capturar relaciones de características significativas. Los autores aseguran que el modelo exhibe una fuerte capacidad predictiva y una codificación de características escalable, ofreciendo un equilibrio mejorado entre la calidad del subtitulado, la eficiencia computacional y la escalabilidad. El trabajo futuro se dirige hacia la incorporación de preentrenamiento multimodal avanzado, la adaptación a conjuntos de datos específicos de dominio y el despliegue de modelos ligeros para aplicaciones de borde en tiempo real.