Enhancing Layer Interaction Using Key-Correlated Layer Attention

Resumen Técnico: Mejora de la Interacción entre Capas mediante la Atención de Capas Correlacionadas por Clave (KCLA)

Planteamiento del Problema
Los avances recientes en el diseño de redes neuronales profundas (DNN) han utilizado la Atención de Capas (Layer Attention) para mejorar las interacciones entre capas, tratando las salidas de las capas como tokens de una secuencia. Sin embargo, la atención de capa estándar sufre de una complejidad computacional cuadrática ($O(D^2)$) respecto a la profundidad de la red $D$, junto con un crecimiento significativo de la memoria debido a la necesidad de almacenar en caché las salidas de todas las capas precedentes.

Para abordar estos problemas, trabajos previos introdujeron la Atención de Capas Recurrente (RLA) y aproximaciones lineales como MRLA-L. Mientras que RLA reduce la sobrecarga computacional mediante el almacenamiento en caché de las representaciones de Clave ($K$) y Valor ($V$), funciona como un mecanismo estático donde las características de las capas tempranas permanecen fijas una vez generadas, lo que limita las actualizaciones dinámicas de información. Además, las aproximaciones lineales como MRLA-L intentan reducir la complejidad mediante la aproximación lineal de la salida de la atención. Los autores observan que este enfoque causa que las contribuciones de las capas anteriores se vean atenuadas exponencialmente (como se muestra en su análisis de los coeficientes de decaimiento $\beta_l$), lo que dificulta el establecimiento de dependencias de largo alcance entre capas efectivas. Adicionalmente, los métodos de linealización existentes a menudo no logran abordar el cuello de botella de memoria del almacenamiento en caché de las salidas de las capas previas o los desafíos específicos de aplicar la atención lineal a estructuras de DNN jerárquicas.

Metodología: Atención de Capas Correlacionada por Clave (KCLA)
La propuesta Atención de Capas Correlacionada por Clave (KCLA) es un novedoso método de atención de capa linealizado diseñado para lograr una complejidad computacional lineal preservando las actualizaciones dinámicas y el modelado de dependencias de largo alcance. La metodología se basa en tres ideas y mecanismos centrales:

1. Observación de la Correlación de Claves: El análisis empírico revela que las representaciones de Clave a través de diferentes capas en la atención de capa estándar exhiben una alta similitud de coseno. Aprovechando esto, KCLA aproxima la salida de la atención de capa estándar factorizando la transformación lineal del Valor ($V$) y aproximando las representaciones de Clave como direccionalmente correlacionadas. Esto permite la agregación de las salidas de las capas en un estado recurrente ($U^t$) sin almacenar el historial completo de las matrices $K$ y $V$.
2. Restauración Dinámica de Softmax: Un desafío crítico al linealizar la atención es la pérdida de la operación softmax, la cual es esencial para normalizar las puntuaciones de atención y permitir ajustes dinámicos de la distribución. KCLA introduce un mecanismo de atención de múltiples cabezales modificado para restaurar esta capacidad:
   • Emplea múltiples cabezales con diferentes factores de escala (temperaturas) derivados de una secuencia logarítmicamente uniforme.
   • Utiliza una transformación lineal aprendible $\Phi$ para mapear la norma de la Consulta ($\|Q^t\|$) y el ángulo de correlación ($\cos \theta$) en pesos específicos de cada cabezal.
   • Esto permite al modelo ajustar dinámicamente la distribución de atención en cada capa, simulando efectivamente el papel del parámetro de temperatura de la función softmax sin calcular la matriz completa de $N \times N$.
3. Adaptación entre Etapas (Cross-Stage Adaptation): Para acomodar DNN con tamaños de mapa de características y dimensiones de canal variables a través de las etapas, KCLA implementa un procedimiento de adaptación específico (Algoritmo 1). Esto implica normalizar las características acumuladas, aplicar una convolución $1 \times 1$ con salto (stride) para el submuestreo y recomputar los factores de normalización para cerrar las disparidades dimensionales entre etapas.

Contribuciones Clave

• Observación y Demostración: Los autores demuestran que las representaciones de Clave a través de diferentes capas en la atención de capa estándar están altamente correlacionadas en dirección, proporcionando la base teórica para la aproximación lineal.
• Nueva Linealización: Se propone KCLA como un método de atención de capa linealizado que evita la naturaleza estática de RLA y la atenuación exponencial de MRLA-L. Logra una complejidad computacional lineal y una complejidad espacial fija independiente de la profundidad de la red.
• Aproximación de Softmax Dinámica: Se introduce un nuevo método para aproximar la operación softmax utilizando el escalado de múltiples cabezales y transformaciones aprendibles, permitiendo el ajuste dinámico y adaptativo de las distribuciones de atención en cada capa.

Resultados Experimentales
Los autores evaluaron KCLA en tres tareas de visión fundamentales: clasificación de imágenes, detección de objetos y segmentación de imágenes médicas.

• Clasificación de Imágenes: En CIFAR-100, KCLA logró el mejor rendimiento a través de profundidades de ResNet de 20, 56 y 110, superando a MRLA-L, DLA-L y la atención de capa estándar. En ImageNet-1K, la integración de KCLA en ResNet-50 y ResNet-101 mejoró la precisión Top-1 en un +1.7% y +1.4%, respectivamente, respecto al baseline. Mantuvo una ventaja ligera pero consistente sobre MRLA-L utilizando significativamente menos parámetros que el método dinámico DLA-L.
• Detección de Objetos: En el conjunto de datos COCO2017 utilizando Faster R-CNN y Mask R-CNN, KCLA demostró un rendimiento competitivo. Con ResNet-101, superó a la arquitectura más compleja DLA-L en varias métricas (AP$_{75}$, AP$_S$, AP$_M$) utilizando 2.2 millones de parámetros menos.
• Segmentación de Imágenes Médicas: En los conjuntos de datos ISIC2017, ISIC2018 y Kvasir-SEG, KCLA superó consistentemente a MRLA-L y al baseline ResNet-UNet, logrando puntuaciones de mIoU y DSC superiores con una sobrecarga mínima de parámetros.
• Análisis de Complejidad: A diferencia de MRLA-B, cuya memoria y FLOPs crecen cuadráticamente con la profundidad, KCLA exhibe un crecimiento lineal tanto en la carga computacional como en el costo de memoria.

Significancia y Reivindicaciones
El artículo afirma que KCLA logra cerrar la brecha entre la eficiencia de la atención lineal y el poder de representación de la atención estándar. Al aproximar directamente el mecanismo de atención de capa estándar basándose en las correlaciones de Clave, KCLA evita los problemas de representación de características estáticas de RLA y la degradación de la dependencia de largo alcance de MRLA-L.

Los autores enfatizan que su método mantiene la capacidad de modelar dependencias de largo alcance entre capas de manera eficiente, ofreciendo un equilibrio deseable entre eficiencia y rendimiento. Señalan que el método propuesto para restaurar la función softmax es un factor crítico para lograr estos resultados. Aunque la implementación actual depende de parámetros de temperatura preestablecidos, los autores identifican el diseño de una atención de capa de múltiples cabezales con temperatura dinámica y totalmente auto-adaptativa como una dirección para la optimización futura. El código está disponible públicamente para facilitar la reproducibilidad.