Una imagen vale más que 16x16 palabras: Transformers para el reconocimiento de imágenes a escala

Guía de lectura

Problema y motivación

Las arquitecturas convolucionales (p. ej., ResNet) han sido durante mucho tiempo la base de la visión por computadora moderna debido a sesgos inductivos incorporados como la equivariancia a traslaciones y la localidad bidimensional. Aunque los Transformers se convirtieron en el estándar en el procesamiento del lenguaje natural (NLP) gracias a su eficiencia computacional y su enorme escalabilidad, los intentos previos de trasladar la autoatención a la visión solían acoplar la atención con backbones de CNN o utilizaban patrones de atención local/axial especializados que resultaban difíciles de escalar eficazmente en aceleradores de hardware. Este artículo investiga si una arquitectura Transformer pura, aplicada directamente a secuencias de parches de imagen con modificaciones mínimas, puede rivalizar o superar a las CNN del estado del arte.

Ideas clave

• Imágenes como secuencias de palabras: Una imagen $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$ se divide en parches sin solapamiento de tamaño $P\times P$. Los parches aplanados se tratan de manera idéntica a los tokens de NLP.
• Sesgo inductivo mínimo: A diferencia de las CNN, donde la localidad espacial y la equivariancia a traslaciones están presentes en cada capa, ViT utiliza autoatención global entre todos los parches. La estructura 2D solo se toma en cuenta al extraer los parches y al interpolar los embeddings de posición para el ajuste fino (fine-tuning) a mayor resolución.
• La escala supera al sesgo inductivo: Sin un sesgo inductivo fuerte, ViT rinde por debajo de ResNet en conjuntos de datos pequeños o medianos (como ImageNet estándar sin una regularización estricta). Sin embargo, cuando se preentrena en conjuntos de datos masivos (ImageNet-21k con 14 millones de imágenes o JFT-300M con 303 millones de imágenes), ViT iguala o supera el rendimiento de las CNN consumiendo sustancialmente menos FLOPs de entrenamiento.

Método

1. Embedding de parches y construcción de la entrada: Dado un tamaño de parche $(P,P)$, una imagen produce $N=HW/P^2$ parches. Cada parche aplanado ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^2\cdot C}$ se proyecta linealmente mediante una matriz $\mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D}$ a la dimensión $D$. Se antepone un token de clasificación aprendible ${\mathbf{x}}_{\text{class}}$ y se añaden embeddings de posición aprendibles 1D ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}$: $${\mathbf{z}}_0=[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos}$$
2. Capas del codificador Transformer: La secuencia pasa a través de $L$ bloques idénticos que alternan Multihead Self-Attention (MSA) y bloques MLP, empleando pre-LayerNorm (LN) y conexiones residuales: $${\mathbf{z}}_{\ell }^{\prime }=\mathrm{MSA}(\mathrm{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -1},\ell =1\dots L$$ $${\mathbf{z}}_{\ell }=\mathrm{MLP}(\mathrm{LN}({\mathbf{z}}_{\ell }^{\prime }))+{\mathbf{z}}_{\ell }^{\prime },\ell =1\dots L$$
3. Representación y cabezal: La representación se obtiene a partir del primer token de salida: $$\mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)$$ Un pequeño MLP (con una capa oculta) actúa como cabezal de clasificación durante el preentrenamiento, el cual se sustituye por una única proyección lineal ($D\times K$) durante el ajuste fino.
4. Ajuste fino a mayor resolución: Al realizar fine-tuning con resoluciones de imagen más altas (p. ej., $384\times 384$ o $512\times 512$), el tamaño de parche $P$ se mantiene constante, lo que incrementa la longitud de la secuencia $N$. Los embeddings de posición 1D preentrenados se ajustan mediante interpolación bicúbica 2D en función de sus coordenadas espaciales originales en la cuadrícula.

Resultados

• Comparación con el estado del arte (SOTA):
  • ViT-H/14 preentrenado en JFT-300M alcanzó un 88.55% en ImageNet, 90.72% en ImageNet ReaL, 94.55% en CIFAR-100 y 77.63% en la suite de 19 tareas de VTAB, requiriendo 2.5k días-núcleo de TPUv3.
  • ViT-L/16 preentrenado en JFT-300M logró un 87.76% en ImageNet con solo 0.68k días-núcleo de TPUv3, superando a BiT-L (ResNet152x4, 87.54%, 9.9k días-núcleo de TPUv3) y aproximándose a Noisy Student EfficientNet-L2 (88.4/88.5%, 12.3k días-núcleo de TPUv3).
  • Preentrenamiento en conjuntos de datos públicos: ViT-L/16 preentrenado en ImageNet-21k alcanzó un 85.30% en ImageNet.
• Compromiso entre cómputo y rendimiento: A lo largo de los experimentos de escalabilidad en JFT-300M, ViT puro dominó a las ResNet en la frontera de Pareto, requiriendo de $2\times$ a $4\times$ menos cómputo para lograr una precisión equivalente.
• Representaciones internas: Las capas inferiores contienen cabezales de atención que atienden a toda la imagen junto con otros con distancias de atención cortas y localizadas. Por lo general, la distancia de atención aumenta de manera monótona con la profundidad de la red.

Limitaciones y preguntas abiertas

• Demanda de datos: ViT generaliza deficientemente en conjuntos de datos reducidos a menos que se disponga de un preentrenamiento a gran escala (ImageNet-21k o JFT-300M) o de una fuerte regularización.
• Brecha en el aprendizaje autosupervisado: Los experimentos iniciales con predicción de parches enmascarados alcanzaron un 79.9% en ImageNet (ViT-B/16), lo que dejó una brecha de rendimiento del 4% en comparación con el preentrenamiento supervisado.
• Aplicación a tareas de predicción densa: El trabajo inicial se enfoca exclusivamente en la clasificación; adaptar los Transformers basados puramente en parches a tareas espaciales densas, como la detección de objetos y la segmentación semántica, continúa siendo un desafío de investigación activo.