Layer by layer, module by module: Choose both for optimal OOD probing of ViT

Este estudio demuestra que el rendimiento óptimo en la detección de cambios de distribución en Vision Transformers se logra combinando el análisis de capas intermedias y módulos específicos, donde la red feedforward es superior ante grandes desplazamientos de distribución y la atención multi-cabeza normalizada lo es ante desplazamientos débiles.

Lo esencial

¡Hola! Imagina que tienes un chef de cocina de clase mundial (el modelo de Inteligencia Artificial) que ha pasado años cocinando en una cocina muy específica, usando ingredientes de alta calidad y recetas perfectas. Este chef es un "modelo fundacional" (como un Vision Transformer o ViT) que ha aprendido a reconocer cosas increíblemente bien... pero solo en su cocina original.

El problema surge cuando le pides que cocine en una fiesta al aire libre con ingredientes diferentes, mal iluminados o con el viento molestando (esto es lo que los expertos llaman "datos fuera de distribución" o OOD).

Este paper es como un estudio que investiga dónde, exactamente, dentro de la mente del chef, se encuentra la mejor receta para adaptar su cocina a esta fiesta desordenada.

Aquí tienes la explicación sencilla, paso a paso:

1. El mito del "Final Perfecto"

Antes de este estudio, todos pensaban que la mejor parte de la mente del chef era su último pensamiento (la capa final del modelo). La lógica era: "Si el chef ha pasado por 12 pasos de preparación, el paso 12 debe ser el más perfecto".

• La analogía: Imagina que el chef está preparando un pastel. Pensábamos que el momento en que pone la última capa de glaseado (la capa final) era cuando el pastel estaba mejor.
• El descubrimiento: El estudio descubrió que, si los ingredientes de la fiesta son muy diferentes a los de su cocina original, el glaseado final arruina todo. El chef se vuelve demasiado rígido y específico para su cocina original, perdiendo la flexibilidad necesaria para la fiesta.

2. La solución: ¡Mira a mitad de camino!

Los investigadores probaron "escuchar" al chef en diferentes momentos de su proceso de pensamiento, no solo al final.

• La analogía: En lugar de preguntar "¿Qué pastel vas a servir?" al final, preguntaron: "¿Qué estás pensando mientras cortas las fresas?" o "¿Qué sientes cuando mezclas la masa?".
• El hallazgo: Cuando el entorno cambia mucho (lluvia, viento, ingredientes raros), las capas intermedias (el pensamiento a mitad de camino) son mucho más inteligentes y adaptables que el final. Son como un chef que aún recuerda cómo adaptar la receta básica, mientras que el final se ha "atascado" en la receta original.

3. No todos los pasos son iguales (El detalle fino)

Aquí es donde el estudio se vuelve muy interesante. Dentro de cada "paso" de pensamiento del chef, hay varias acciones:

1. Mirar los ingredientes (Atención).
2. Mezclarlos (Red neuronal).
3. Sazonar (Normalización).
4. Darle un toque final (Salida).

El estudio descubrió que, si el entorno es muy difícil (lluvia fuerte, ingredientes extraños):

• Lo mejor: Preguntar al chef justo después de que mezcla los ingredientes (la activación dentro de la red neuronal, llamada "Feedforward"). Es como preguntar: "¿Qué sabor tiene la masa ahora mismo?". Esta es la información más pura y útil.
• Lo peor: Preguntar después de que el chef comprime o resume la información (la capa final de la mezcla). Aquí es donde pierde detalles importantes.
• El caso especial: Si la fiesta es tranquila y los ingredientes son casi los mismos de siempre (entorno "In-Distribution"), entonces sí, el final (el glaseado) es perfecto.

4. ¿Qué nos enseña esto para el futuro?

El papel nos da dos reglas de oro para cuando usamos estas inteligencias artificiales en el mundo real:

1. Si todo va bien (entorno estable): Usa la capa final del modelo. Es la más precisa.
2. Si hay cambios o sorpresas (entorno inestable): No uses la capa final. En su lugar, "escucha" al modelo en medio del proceso, específicamente justo después de que procesa la información (la activación).

En resumen

Imagina que el modelo es un viajero.

• Si viaja por la misma carretera de siempre, el destino final (la capa final) es donde tiene la mejor vista.
• Pero si el camino cambia, hay baches y el clima es malo, el viajero tiene la mejor visión y adaptabilidad mientras está caminando por el medio del bosque (las capas intermedias), antes de llegar a un destino que ya no existe.

La lección: No confíes ciegamente en el final del proceso. A veces, la respuesta más inteligente está en el medio del camino, especialmente cuando el mundo cambia.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo en español, estructurado según los puntos solicitados:

Título: Capa por Capa, Módulo por Módulo: Elegir Ambos para la Sonda Óptima de ViT en Distribución Fuera de la Original (OOD)

1. Problema

Los modelos fundacionales, específicamente los Transformadores de Visión (ViT), a menudo sufren degradación de rendimiento cuando se despliegan en entornos reales que presentan desplazamientos de distribución (distribution shifts) respecto a sus datos de preentrenamiento.

• Hipótesis convencional: Se asumía que las representaciones de la capa final del modelo eran las más óptimas para tareas de clasificación.
• Desafío actual: Estudios recientes sugieren que las capas intermedias pueden ser superiores, pero los resultados son contradictorios dependiendo del tipo de modelo (autoregresivo vs. ViT estándar) y del tipo de datos (in-distribution vs. out-of-distribution).
• Brecha de conocimiento: No está claro cuándo y por qué las capas intermedias superan a las finales en ViTs, ni qué componentes específicos dentro de un bloque transformador (atención, feed-forward, normalización) ofrecen las representaciones más robustas ante cambios de distribución.

2. Metodología

Los autores realizaron un estudio exhaustivo utilizando un ViT preentrenado de 86M parámetros (entrenado en ImageNet-21k) y evaluado mediante sondeo lineal (linear probing) en un conjunto diverso de 11 benchmarks de clasificación de imágenes.

• Configuración Experimental:
  • Datos: Incluyeron datos in-distribution (ID) como CIFAR-10, CIFAR-100, Flowers102 y Pets, y datos out-of-distribution (OOD) con desplazamientos severos (ruido, desenfoque, cambios de dominio como Clipart y Sketch de DomainNet).
  • Protocolo: Compararon el rendimiento del modelo congelado (solo preentrenado) frente al modelo ajustado (finetuned) en cada dataset.
  • Análisis Granular: En lugar de sondear solo la salida del bloque transformador (residual connection 2, RC2), analizaron individualmente los 8 componentes operativos dentro de cada capa:
    1. LN1 (Normalización antes de la Atención).
    2. MHA (Salida del Módulo de Atención Multi-cabeza).
    3. RC1 (Conexión residual tras la atención).
    4. LN2 (Normalización antes de la red Feed-Forward).
    5. FC1 (Primera capa totalmente conectada).
    6. Act (Activación GeLU entre FC1 y FC2).
    7. FC2 (Segunda capa totalmente conectada).
    8. RC2 (Salida final del bloque).

3. Contribuciones Clave

El artículo aporta dos hallazgos principales que desafían el conocimiento establecido:

1. El Desplazamiento de Distribución es la Causa Raíz: Se demuestra que la degradación del rendimiento en las capas finales no es un defecto inherente al entrenamiento supervisado o auto-supervisado, sino una consecuencia directa del desplazamiento de distribución entre los datos de preentrenamiento y los datos de la tarea downstream.

   • En escenarios ID (sin desplazamiento), las capas finales siempre son superiores.
   • En escenarios OOD (con desplazamiento), las capas intermedias son significativamente más robustas.

2. Optimalidad del Módulo Específico: Se revela que la elección del punto de sondeo dentro de la arquitectura es crítica y depende de la magnitud del desplazamiento:

   • Desplazamiento Fuerte (OOD): La activación dentro de la red Feed-Forward (Act, después de GeLU) en las capas intermedias ofrece el mejor rendimiento.
   • Desplazamiento Débil o Nulo (ID): La salida normalizada del módulo de atención (LN2) o la salida del bloque (RC2) son óptimas.
   • Hallazgo Negativo: La salida de la segunda capa totalmente conectada (FC2) es consistentemente la peor opción para sondear, especialmente bajo desplazamientos fuertes.

4. Resultados

• Comportamiento por Capas: En datasets con alto desplazamiento (ej. "Speckle Noise", "Sketch"), la precisión de las capas finales del modelo congelado cae drásticamente, mientras que las capas intermedias mantienen una precisión más estable y alta. En contraste, en datasets ID (ej. Flowers102), la precisión aumenta monótonamente hasta la última capa.
• Comparación de Módulos (Tabla 1):
  • Act (Activación): Logró la mayor tasa de victorias (mejor precisión) en la mayoría de los datasets OOD, superando a los componentes estándar por un margen considerable en casos de ruido severo.
  • FC2: Fue el peor componente en 10 de los 12 datasets evaluados, sugiriendo que la compresión de dimensiones en la segunda capa del Feed-Forward reduce la separabilidad lineal bajo distribución cambiante.
  • LN2: Mostró un rendimiento estable y robusto, siendo una opción segura si el grado de desplazamiento es difícil de detectar.
• Interpretación: Los autores sugieren que FC1 y Act operan en un espacio de mayor dimensión (4d), lo que promueve el "desenredo" de características (feature disentanglement) y filtra el ruido, mientras que FC2 comprime la información, perdiendo separabilidad lineal.

5. Significado e Impacto

Este trabajo tiene implicaciones prácticas importantes para el despliegue de modelos de visión por computadora en el mundo real:

• Guía para la Detección de Desplazamiento: Proporciona una métrica basada en la arquitectura para inferir si un modelo está operando bajo una distribución conocida o desconocida, simplemente observando qué capa o módulo ofrece la mejor precisión de sondeo.
• Mejora de la Robustez: Ofrece una estrategia de bajo costo computacional (sin necesidad de finetuning completo) para mejorar la precisión en entornos OOD: sondear la activación de la red Feed-Forward en capas intermedias.
• Refinamiento de la Arquitectura: Sugiere que la práctica estándar de usar la salida del bloque transformador (RC2) como representación final no es óptima para todos los escenarios, abriendo nuevas direcciones para el análisis de representaciones ocultas en transformadores mediante medidas teóricas de la información y geometría.

En resumen, el papel concluye que para obtener la máxima robustez en ViTs frente a cambios de distribución, no basta con elegir la capa correcta; es fundamental elegir el módulo correcto (específicamente la activación intermedia de la red Feed-Forward) dentro de esa capa.