An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Guide de lecture

Problème et motivation

Les architectures convolutives (ex. ResNet) constituent depuis longtemps le fondement de la vision par ordinateur moderne en raison de biais inductifs codés en dur tels que l'équivariance par translation et la localité bidimensionnelle. Alors que les Transformers sont devenus le standard dans le traitement du langage naturel (NLP) grâce à leur efficacité computationnelle et leur passage à l'échelle massif, les tentatives antérieures visant à adapter la self-attention à la vision associaient généralement l'attention à des backbones CNN ou utilisaient des motifs d'attention locale/axiale spécialisés, difficiles à accélérer efficacement sur le matériel dédié. Cet article examine si une architecture Transformer pure, appliquée directement à des séquences de patchs d'images avec un minimum de modifications, peut rivaliser avec les CNN de l'état de l'art, voire les surpasser.

Idées clés

• Les images comme séquences de mots : Une image $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$ est découpée en patchs sans chevauchement de taille $P\times P$. Les patchs aplatis sont traités de manière identique à des tokens en NLP.
• Biais inductif minimal : Contrairement aux CNN où la localité spatiale et l'équivariance par translation sont intégrées à chaque couche, ViT utilise une self-attention globale sur l'ensemble des patchs. La structure 2D n'intervient que lors de l'extraction des patchs et lors de l'interpolation des position embeddings pour le fine-tuning à plus haute résolution.
• L'échelle l'emporte sur le biais inductif : Dépourvu de fort biais inductif, ViT sous-performe face aux ResNets sur des jeux de données de taille modeste ou moyenne (comme ImageNet standard sans forte régularisation). En revanche, lorsqu'il est pré-entraîné sur de grands jeux de données (ImageNet-21k avec 14M d'images ou JFT-300M avec 303M d'images), ViT égale ou dépasse les performances des CNN tout en consommant nettement moins de FLOPs à l'entraînement.

Méthode

1. Patch Embedding et construction de l'entrée : Pour une taille de patch donnée $(P,P)$, une image produit $N=HW/P^2$ patchs. Chaque patch aplati ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^2\cdot C}$ est projeté linéairement par une matrice $\mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D}$ vers une dimension $D$. Un token de classification apprenable ${\mathbf{x}}_{\text{class}}$ est préfixé, et des position embeddings 1D apprenables ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}$ sont ajoutés : $${\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. Couches d'encodeur Transformer : La séquence traverse $L$ blocs identiques alternant Multihead Self-Attention (MSA) et blocs MLP, avec pre-LayerNorm (LN) et connexions résiduelles : $${\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. Représentation et tête de classification : La représentation est extraite du premier token de sortie : $$\mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)$$ Un petit MLP (avec une couche cachée) sert de tête de classification durant le pré-entraînement, remplacé par une simple projection linéaire ($D\times K$) lors du fine-tuning.
4. Fine-tuning à plus haute résolution : Lors du fine-tuning à des résolutions d'image supérieures (ex. $384\times 384$ ou $512\times 512$), la taille de patch $P$ reste fixe, ce qui augmente la longueur de séquence $N$. Les position embeddings 1D pré-entraînés sont alors ajustés par interpolation bicubique 2D en fonction de leurs coordonnées d'origine sur la grille spatiale.

Résultats

• Comparaison à l'état de l'art :
  • ViT-H/14 pré-entraîné sur JFT-300M a atteint 88,55 % sur ImageNet, 90,72 % sur ImageNet ReaL, 94,55 % sur CIFAR-100 et 77,63 % sur la suite de 19 tâches VTAB, en nécessitant 2,5k TPUv3-core-days.
  • ViT-L/16 pré-entraîné sur JFT-300M a atteint 87,76 % sur ImageNet avec seulement 0,68k TPUv3-core-days, surpassant BiT-L (ResNet152x4, 87,54 %, 9,9k TPUv3-core-days) et s'approchant de Noisy Student EfficientNet-L2 (88,4/88,5 %, 12,3k TPUv3-core-days).
  • Pré-entraînement sur jeux de données publics : ViT-L/16 pré-entraîné sur ImageNet-21k a atteint 85,30 % sur ImageNet.
• Compromis calcul / performance : À travers les expériences d'échelle sur JFT-300M, le ViT pur a dominé les ResNets sur le front de Pareto, exigeant 2 à 4 fois moins de calcul pour atteindre une précision équivalente.
• Représentations internes : Les couches inférieures contiennent des têtes d'attention qui traitent l'image entière aux côtés d'autres têtes avec des distances d'attention courtes et localisées. La distance d'attention augmente généralement de façon monotone avec la profondeur.

Limites et questions ouvertes

• Besoin massif de données : ViT généralise mal sur de petits jeux de données sans données massives de pré-entraînement (ImageNet-21k ou JFT-300M) ou sans forte régularisation.
• Écart en auto-supervisé : Les premières expériences de prédiction de patchs masqués ont atteint 79,9 % sur ImageNet (ViT-B/16), laissant un écart de performance de 4 % par rapport au pré-entraînement supervisé.
• Application à la prédiction dense : Ces travaux initiaux se concentrent exclusivement sur la classification ; adapter des architectures transformers purement à base de patchs à des tâches spatiales denses (comme la détection d'objets ou la segmentation sémantique) reste un sujet de recherche actif.