LAYA: Layer-wise Attention Aggregation for Interpretable Depth-Aware Neural Networks

Resumen Técnico: LAYA (Agregación de Atención por Capas)

Planteamiento del Problema
Las redes neuronales profundas (DNN) suelen depender exclusivamente de la representación producida por la última capa oculta ($h_L$) para generar predicciones. Este diseño asume implícitamente que el embedding más profundo encapsula plenamente toda la información semántica relevante extraída mediante las transformaciones precedentes. Sin embargo, la evidencia empírica sugiere que las capas intermedias contienen información rica y complementaria que abarca desde patrones de bajo nivel hasta abstracciones de alto nivel. Al descartar estas representaciones intermedias en la etapa de salida, las arquitecturas estándar no solo pueden perder información discriminativa, sino también oscurecer la contribución de los diferentes niveles de abstracción a la decisión final. Además, los métodos existentes para interpretar estas contribuciones son mayoritariamente post-hoc, lo que requiere un análisis externo después de la inferencia, en lugar de ser intrínsecos al mecanismo de predicción del modelo.

Metodología
El artículo presenta LAYA (Layer-wise Attention Aggregator), un módulo de salida ligero y agnóstico a la arquitectura, diseñado para reemplazar la cabeza de clasificación estándar. En lugar de proyectar únicamente el embedding final, LAYA agrega dinámicamente las representaciones internas de todas las capas ocultas $\{h_i\}_{i=1}^L$.

El mecanismo opera de la siguiente manera:

1. Proyección: Cada representación oculta $h_i$ se mapea en un espacio latente compartido de dimensión $d^*$ mediante un adaptador aprendible $g_i(\cdot)$, resultando en $z_i = g_i(h_i)$.
2. Puntuación: Estas representaciones proyectadas son opcionalmente transformadas por un mapeo compartido $\psi(\cdot)$ (identidad o un MLP ligero) y se concatenan. Un MLP de puntuación procesa esta concatenación para producir puntuaciones de relevancia no normalizadas $s(x) \in \mathbb{R}^L$.
3. Atención: Las puntuaciones se convierten en coeficientes de atención condicionados a la entrada $\alpha_i(x)$ utilizando una función softmax con escala de temperatura:
   $$\alpha_i(x) = \frac{\exp(s_i(x)/\tau)}{\sum_{j=1}^L \exp(s_j(x)/\tau)}$$
   donde $\tau$ controla la nitidez de la distribución.
4. Agregación y Predicción: La representación agregada final se calcula como una suma ponderada $h_{agg} = \sum_{i=1}^L \alpha_i(x) z_i$, la cual se pasa a la cabeza de salida específica para la tarea de clasificación.

Crucialmente, los pesos de atención $\alpha_i(x)$ se generan durante la pasada hacia adelante (forward pass) como parte del proceso de predicción, sirviendo como puntuaciones de atribución de capa intrínsecas sin requerir cálculos de retropropagación adicionales o modelos de explicación externos.

Contribuciones Clave

1. Revisión de la Capa de Salida: El trabajo destaca la limitación de las cabezas de salida convencionales que ignoran las representaciones intermedias y propone un cambio hacia la agregación consciente de la profundidad.
2. Interpretabilidad Intrínseca: LAYA proporciona un mecanismo donde el proceso de decisión del modelo es inherentemente interpretable. Los coeficientes de atención aprendidos cuantifican explícitamente la contribución de cada capa a la predicción final, eliminando la necesidad de herramientas de explicación post-hoc como Integrated Gradients o Grad-CAM para la atribución a nivel de capa.
3. Agnosticismo de Arquitectura: El módulo funciona como un reemplazo directo para las cabezas de salida estándar, dejando la red base (por ejemplo, CNNs, Transformers, MLPs) completamente inalterada.
4. Evaluación Exhaustiva: El artículo presenta una evaluación riguroa a través de tres jerarquías representacionales distintas: un Perceptrón Multicapa (MLP) en Fashion-MNIST, una Red Neuronal Convolucional (CNN) en CIFAR-10 y un Vision Transformer (ViT) congelado en una tarea de clasificación de estilo artístico.

Resultados Experimentales

• Rendimiento Predictivo: LAYA logra un rendimiento predictivo competitivo, y en algunos casos superior, comparado con las líneas base estándar (LastLayer), la concatenación estática (Concat) y la mezcla global (ScalarMix).
  • En Fashion-MNIST y CIFAR-10, LAYA alcanzó la precisión media más alta entre todas las cabezas probadas.
  • En el conjunto de datos Best Artworks, aunque la concatenación estática fue la mejor, LAYA superó significativamente a la línea base estándar LastLayer, demostrando que las representaciones intermedias contienen información valiosa incluso en redes base preentrenadas y congeladas.
• Fidelidad de las Explicaciones: El análisis cuantitativo utilizando Faithfulness (correlación de Spearman entre las puntuaciones de atribución y los efectos de perturbación de exclusión de una capa) y Deletion@1 (caída de confianza al eliminar la capa con mayor rango) muestra que las puntuaciones intrínsecas de LAYA son altamente fieles.
  • En Fashion-MNIST, la puntuación de fidelidad de LAYA (0.8575) se acercó a la de los métodos post-hoc basados en gradientes (aprox. 0.89), superando significativamente a los pesos globales estáticos.
  • En el conjunto de datos Best Artworks, LAYA logró la puntuación Deletion@1 más alta (0.5787), indicando que la capa que identificó como más relevante era, de hecho, la más influyente para la predicción.
• Patrones de Interpretabilidad: El análisis cualitativo reveló una utilización de la profundidad dependiente de la tarea. Por ejemplo, en CIFAR-10, la atención se concentró casi totalmente en la última capa, mientras que en Best Artworks, la atención se distribuyó de forma dispersa a través de bloques específicos de la etapa intermedia y tardía del ViT (por ejemplo, bloques 7, 9, 11, 12), variando según el estilo artístico.

Significancia y Reivindicaciones
El artículo postula que LAYA ofrece una alternativa fundamentada al "colapso" estándar de las representaciones jerárquicas en la etapa de salida. Su principal significancia no radica en dominar universalmente la precisión de la clasificación, sino en su capacidad de combinar un rendimiento predictivo competitivo con una atribución de capa dependiente de la entrada y directamente inspeccionable.

El autor afirma que LAYA:

• Proporciona interpretabilidad intrínseca, haciendo transparente la dependencia del modelo hacia niveles de abstracción específicos sin herramientas externas.
• Revela patrones estructurados y dependientes de la tarea de utilización de la profundidad, mostrando que diferentes tareas e incluso diferentes clases dentro de una tarea dependen de diferentes combinaciones de profundidades de capa.
• Ofrece una base para aplicaciones futuras tales como la compresión de modelos guiada por atención (poda de capas con atención insignificante), mecanismos de salida temprana adaptativos (early-exit) y análisis diagnóstico de fallos del modelo basados en cambios de atención.

El trabajo concluye que, si bien el enfoque es la visión por computadora, el marco es aplicable a otras arquitecturas de secuencias, ofreciendo potencialmente nuevos conocimientos sobre cómo se refina la información a través de la profundidad en modelos como los Transformers.