Vision Transformers Need Better Token Interaction

Resumen Técnico: Los Transformadores de Visión Necesitan una Mejor Interacción entre Tokens

Enunciado del Problema
Los Transformadores de Visión (ViT) son reconocidos por aprender representaciones fuertes a nivel de imagen, sin embargo, sus representaciones a nivel de parche a menudo se degradan durante el entrenamiento prolongado, volviéndose menos efectivas para tareas de predicción densa como la segmentación semántica y la estimación de profundidad. Aunque investigaciones anteriores han atribuido esta "degradación densa" a artefactos de alta norma o tokens atípicos, este artículo argumenta que estos son síntomas y no la causa raíz. Los autores identifican el problema central como la difusión semántica: un atajo de optimización donde la información semántica global se propaga indiscriminadamente a través de los tokens de parche más allá de lo que está justificado localmente. Este fenómeno permite que el modelo mejore el reconocimiento global explotando interacciones de tokens sin restricciones, pero a costa de degradar la estructura espacial requerida para tareas densas.

Metodología
Los autores proponen que la solución no es eliminar el contexto global ni suprimir completamente los tokens de fondo, sino hacer que las interacciones entre tokens sean más selectivas. Introducen la atención dispersa como una intervención mínima para modificar la geometría de la atención.

• Mecanismo: El mecanismo estándar de atención softmax se reemplaza con entmax-1.5. A diferencia de softmax, que produce distribuciones de probabilidad densas, entmax-1.5 es un mapeo de probabilidad disperso que puede asignar pesos cero exactos a interacciones de baja utilidad.
• Implementación: El método preserva la conectividad global de los tokens (a diferencia de la atención de ventana local) y no introduce parámetros adicionales, tokens de registro ni ramas de procesamiento separadas (a diferencia de enfoques de especialización o basados en registros).
• Hipótesis: Al fomentar distribuciones de atención dispersas, el modelo se ve obligado a ser más selectivo en su mezcla de tokens. Esto reduce el "atajo de optimización" de la difusión semántica, evitando que los parches de fondo semánticamente irrelevantes diluyan la estructura local de los parches de primer plano, mientras se mantiene la capacidad de capturar semántica global.

Contribuciones Clave

1. Replanteamiento de la Degradación Densa: El artículo caracteriza la degradación densa no meramente como un problema de artefactos, sino como una consecuencia de la difusión semántica, un atajo de optimización donde las semánticas globales se difunden en los tokens de parche de manera localmente inconsistente.
2. Análisis de los Roles de los Tokens: Mediante el análisis Point-in-Box (PiB) y la sondeo por capas, los autores demuestran que, aunque las capas superficiales mantienen una mejor alineación espacial (PiB más alto), carecen de la riqueza semántica de las capas más profundas. Crucialmente, muestran que el token [CLS] retiene información global complementaria que, al concatenarse con las características de los parches, ayuda significativamente a la predicción densa, lo que sugiere que el contexto global debe preservarse pero gestionarse mejor.
3. Intervención Mínima: El artículo introduce un cambio arquitectónico simple: sustituir softmax por entmax-1.5, que actúa como un sesgo representacional hacia la interacción selectiva de tokens sin requerir supervisión adicional, modelos maestros o modificaciones arquitectónicas complejas.

Resultados Experimentales
El método se evaluó en modelos DINOv1 ViT-S/16 entrenados durante 200 épocas en ImageNet-1K. Los resultados indican que la atención dispersa mejora las representaciones densas sin sacrificar el reconocimiento global:

• Reconocimiento a Nivel de Imagen: La variante Entmax-1.5 logró una precisión Top-1 de 70.06% en el sondeo lineal de ImageNet, superando ligeramente a la línea base CLS (69.50%) y a otras variantes como Register-4 y Especialización.
• Segmentación Semántica: Aquí se observaron las ganancias más significativas. En comparación con la línea base CLS, Entmax-1.5 mejoró:
  • mIoU de VOC: de 42.80 a 48.78 (+5.98).
  • mIoU de ADE20K: de 19.85 a 21.97 (+2.12).
  • mIoU de Cityscapes: de 36.79 a 37.87 (+1.08).
  • Superó tanto a Register-4 como a la especialización CLS/parche en las tres pruebas de referencia de segmentación.
• Estimación de Profundidad: Las mejoras fueron más moderadas pero consistentes en las pruebas de referencia NYUv2, SUN RGB-D y KITTI (reduciendo el RMSE en 0.004–0.005). Los autores señalan que el protocolo de sondeo de profundidad de la línea base ya concatena las características de [CLS] y parche, lo que puede compensar parcialmente la difusión semántica, haciendo más difíciles de lograr las ganancias.
• Análisis Cualitativo: Las visualizaciones PCA de las características de los parches mostraron que Entmax-1.5 produce mapas de características más limpios y espacialmente coherentes, con una mejor separación entre las regiones de primer plano y fondo en comparación con las líneas base.

Significado y Afirmaciones
El artículo afirma que la mezcla selectiva de tokens es un sesgo simple y efectivo para mejorar las representaciones densas de ViT. El significado radica en demostrar que el objetivo de la predicción densa no es eliminar el contexto global, sino refinar cómo interactúan los tokens con él. Al reemplazar softmax con entmax-1.5, los autores muestran que una modificación ligera a la normalización de la atención puede mitigar la difusión semántica, lo que lleva a mejoras sustanciales en la segmentación semántica mientras se mantiene o mejora ligeramente el reconocimiento a nivel de imagen.

Los autores mantienen modestia respecto a sus afirmaciones, señalando que este estudio se limita a DINOv1 ViT-S/16 y aún no ha evaluado backbones más grandes ni programas de preentrenamiento más largos. Posicionan la atención dispersa no como una arquitectura final definitiva, sino como evidencia que respalda el principio más amplio de que hacer que las interacciones entre tokens sean más selectivas es un camino viable para resolver la degradación densa.