Thicker and Quicker: A Jumbo Token for Fast Plain Vision Transformers

Ce papier présente Jumbo, une architecture de Vision Transformer (ViT) qui accélère le traitement en réduisant la largeur des tokens de patchs tout en ajoutant un token global « Jumbo » plus large et efficace, permettant d'améliorer la précision et l'adaptabilité sur diverses tâches sans sacrifier la compatibilité ni la simplicité des ViT standards.

L'essentiel

🚀 Jumbo : Le Super-Héros des "Yeux Numériques"

Imaginez que vous essayez de faire comprendre à un ordinateur ce qu'il voit sur une photo. Pour cela, on utilise des modèles appelés Transformeurs de Vision (ViT). C'est comme si l'ordinateur découpait la photo en milliers de petits morceaux (des "patchs") et les lisait un par un pour comprendre l'image globale.

Le problème ? Ces modèles sont souvent lents et gourmands en énergie. Pour les rendre plus rapides, les ingénieurs ont souvent deux choix difficiles :

1. Rendre le modèle plus petit (comme passer d'un camion à une voiture de sport), mais il devient moins intelligent.
2. Mélanger des technologies différentes (comme ajouter un moteur diesel à une voiture électrique), ce qui le rend complexe et moins flexible.

L'équipe derrière Jumbo a trouvé une troisième voie : rendre le modèle plus gros, mais plus malin, sans le ralentir.

---

🧩 L'Analogie de l'Équipe de Construction

Pour comprendre comment Jumbo fonctionne, imaginons une équipe de construction qui doit rénover une maison (l'image).

1. La méthode classique (ViT normal)

Dans un modèle standard, l'équipe est composée de 196 maçons (les petits morceaux de l'image) et 1 chef de chantier (le token "CLS").

• Les 196 maçons travaillent dur sur les détails (les briques, les fenêtres).
• Le chef de chantier essaie de tout résumer, mais il n'est qu'une seule personne. Il est débordé et ne peut pas tout voir clairement.
• Si on veut aller plus vite, on réduit le nombre de maçons, mais la maison est moins bien faite.

2. La méthode "Jumbo" (Le nouveau héros)

L'équipe de Jumbo dit : "Et si on avait un Super-Chef de Chantier ?"

Au lieu d'avoir un seul chef normal, ils ajoutent un Jumbo :

• C'est un géant : Ce chef est 6 fois plus gros (plus large) que les maçons. Il a une capacité de réflexion énorme.
• Il est spécial : Avant de donner ses ordres, ce géant se divise en 6 petits chefs temporaires pour discuter avec les maçons (c'est l'attention). Ensuite, ils se réunissent pour former le grand chef Jumbo qui prend la décision finale.
• Il est économe : Même s'il est gros, il ne coûte pas cher à l'entretien car il partage ses connaissances avec toutes les étapes de la construction (partage de paramètres).

Le résultat ? Les maçons (les détails) travaillent aussi vite qu'avant, mais le chef (la compréhension globale) est beaucoup plus intelligent. L'équipe finit le travail plus vite et avec une meilleure qualité.

---

🌟 Pourquoi Jumbo est une révolution ?

Voici les 3 super-pouvoirs de Jumbo expliqués simplement :

1. Il est "Universel" (Comme un couteau suisse)

Beaucoup de modèles rapides sont comme des couteaux spécialisés : excellents pour couper du pain, mais inutiles pour ouvrir une bouteille.

• Les modèles rapides actuels (comme EfficientViT) sont souvent trop complexes pour être utilisés dans d'autres domaines (comme la vidéo, les séries temporelles ou l'IA générative).
• Jumbo, lui, reste un couteau suisse. Il garde la simplicité du modèle original. Vous pouvez l'utiliser pour des photos, des vidéos, des données financières ou même du texte, sans avoir à le réinventer.

2. Il est "Écologique" (Moins de gaspillage)

Même si le chef Jumbo est plus gros, il ne consomme pas plus de carburant.

• L'article montre que Jumbo est plus rapide que les meilleurs modèles spécialisés actuels.
• Il permet de faire des tâches complexes (comme la segmentation d'images ou l'apprentissage sans étiquettes) avec moins de puissance de calcul. C'est comme avoir une Ferrari qui consomme moins d'essence qu'une citadine.

3. Il résout le problème du "Chef débordé"

Les chercheurs ont remarqué que dans les modèles classiques, le chef de chantier (le token global) était souvent noyé sous les détails et créait des "artefacts" (des erreurs de compréhension).

• Avec Jumbo, ce chef a assez de "cerveau" pour trier les informations importantes et ignorer le bruit. C'est comme si le chef avait des lunettes de vision nocturne et un casque anti-bruit : il voit tout, clairement.

---

📊 Les Résultats en Bref

L'équipe a testé Jumbo partout :

• Reconnaissance d'images : Il bat les records de précision tout en étant plus rapide.
• Séries temporelles (données qui évoluent dans le temps) : Il est le numéro 1 sur 20 benchmarks différents.
• Robustesse : Même si l'image est floue, bruitée ou corrompue, Jumbo comprend mieux que les autres.
• Apprentissage : Il apprend mieux avec moins d'exemples (ce qu'on appelle l'apprentissage auto-supervisé).

🏁 En Résumé

Jumbo, c'est comme donner un cerveau de génie à un athlète de sprint.
Au lieu de courir plus vite en étant plus petit (ce qui limite sa force), on lui donne un corps plus puissant et plus intelligent, tout en gardant sa légèreté.

C'est une avancée majeure car cela permet d'avoir des modèles à la fois rapides, précis et flexibles, capables de s'adapter à n'importe quelle tâche, du diagnostic médical à la conduite autonome, sans avoir besoin de construire une nouvelle machine pour chaque problème.

Le code et les modèles sont déjà disponibles pour que tout le monde puisse utiliser cette technologie ! 🚀

Résumé technique

1. Problématique

Les Vision Transformers (ViT) sont devenus l'architecture de référence en vision par ordinateur grâce à leur généralité et leur précision. Cependant, ils souffrent d'un manque d'efficacité computationnelle par rapport aux architectures spécialisées (comme les CNNs optimisés ou les modèles hybrides).

Les méthodes existantes pour accélérer les ViT présentent deux inconvénients majeurs :

1. Architectures hybrides : Elles introduisent des convolutions ou des hiérarchies, ce qui brise la simplicité du "Plain ViT" (attention-only, non hiérarchique) et les rend incompatibles avec de nombreuses techniques avancées (apprentissage auto-supervisé par masquage, adaptation au test, traitement de données non-2D).
2. Réduction de la largeur (Shrinking) : Réduire uniformément la largeur des tokens (canaux) pour gagner en vitesse entraîne une perte significative de précision.

Le défi consiste donc à créer une architecture plus rapide et plus précise tout en conservant les propriétés fondamentales du ViT "plain" (non hiérarchique, uniquement basé sur l'attention) pour garantir sa compatibilité avec l'écosystème existant.

2. Méthodologie : L'Architecture "Jumbo"

Les auteurs proposent une nouvelle architecture, Jumbo, qui modifie l'équilibre des capacités de calcul au sein du ViT sans en changer la structure globale.

• Token Jumbo Asymétrique : Au lieu d'avoir une largeur uniforme pour tous les tokens, Jumbo introduit un token global unique (remplaçant le token CLS classique) qui est $J$ fois plus large que les tokens de patchs (par exemple, $J=6$).
• Traitement du Token Jumbo :
  • Avant l'attention : Le token Jumbo large est divisé en $J$ tokens de taille standard pour participer à l'attention multi-têtes (MHSA) avec les tokens de patchs.
  • Après l'attention : Les $J$ tokens sont réassemblés (concaténés) pour reformer le token Jumbo unique.
  • Réseau de Neurones (FFN) : Le token Jumbo est traité par son propre FFN (Feed-Forward Network) élargi, distinct de celui des patchs. Ce FFN est conçu pour augmenter la capacité du modèle à traiter les informations globales.
• Partage de Paramètres (Layer Sharing) : Pour limiter l'augmentation de la mémoire, les paramètres du FFN du token Jumbo sont partagés entre toutes les couches du modèle. Cela réduit considérablement le nombre total de paramètres tout en agissant comme une forme de régularisation.
• Préservation des Propriétés ViT : L'architecture reste non hiérarchique et uniquement basée sur l'attention. Elle conserve donc la capacité de :
  • Gérer des formes d'entrée variées (séries temporelles, vidéo, 3D).
  • Utiliser le "Token Dropping" (abandon de tokens) pour l'entraînement et l'inférence efficaces.
  • Être compatible avec les algorithmes d'apprentissage auto-supervisé (SSL) comme MAE.

3. Contributions Clés

1. Architecture Efficace et Compatible : Jumbo est la première architecture "plain ViT" (attention-only, non hiérarchique) qui surpasse les architectures computationnellement efficaces spécialisées (comme EfficientViT ou MobileNetV4) tout en restant compatible avec l'écosystème ViT.
2. Augmentation Asymétrique de la Capacité : En augmentant la largeur uniquement du token global et en partageant son FFN, le modèle gagne en capacité de représentation sans pénalité de vitesse significative.
3. Généralité Multi-Modalité : La méthode démontre une efficacité supérieure non seulement sur la classification d'images, mais aussi sur la segmentation, l'apprentissage auto-supervisé, l'adaptation au test (TTA) et la modélisation de séries temporelles.
4. Réduction des Artefacts : Comme les "Registers" (tokens d'enregistrement), le token Jumbo aide à réduire les tokens à forte norme (outliers) qui causent des artefacts dans les cartes d'attention, améliorant ainsi la robustesse.

4. Résultats Expérimentaux

Les expériences ont été menées sur plusieurs benchmarks et tâches :

• Classification d'images (ImageNet-1K & 21K) :

  • Jumbo atteint la frontière de Pareto (meilleur compromis vitesse/précision) par rapport aux architectures spécialisées.
  • Sur ImageNet-1K, il surpasse les ViT avec Registers de 0,1 % à 13 % (selon l'échelle, avec des gains massifs sur les modèles Nano et Tiny).
  • Sur ImageNet-21K, les gains augmentent avec la complexité de la tâche (jusqu'à +3,1 % sur ViT-Small), confirmant l'hypothèse que Jumbo bénéficie davantage aux tâches à haute dimensionnalité de sortie.
  • Il est 1,9 fois plus rapide que les ViT+Registers pour une précision équivalente sur ImageNet-21K.

• Segmentation Sémantique (ADE20K) :

  • Jumbo améliore le mIoU de 1,9 % à 3,1 % par rapport aux Registers, en utilisant une tête de segmentation standard.

• Apprentissage Auto-Supervisé (MAE) :

  • Un ViT-Base+Jumbo pré-entraîné avec MAE atteint la même précision qu'un ViT-Large standard, mais avec 2,3 fois moins de paramètres, 3,5 fois moins de FLOPs et 3,1 fois plus de débit.

• Robustesse et Adaptation au Test (ImageNet-C) :

  • Jumbo est plus robuste aux corruptions d'images. Avec l'adaptation au test (SAR), il améliore la précision de +5,2 % par rapport aux Registers.

• Séries Temporelles :

  • Appliqué à PatchTST (transformer pour séries temporelles), Jumbo se classe premier sur 20 benchmarks (univariés et multivariés), surpassant les baselines et les versions avec Registers.

5. Signification et Impact

L'article "Thicker and Quicker" démontre qu'il n'est pas nécessaire de sacrifier la simplicité et la flexibilité des ViT pour obtenir de l'efficacité. En rééquilibrant la capacité computationnelle vers un token global élargi et partagé, les auteurs parviennent à :

• Dépasser les architectures spécialisées (CNNs, modèles hybrides) en termes de compromis vitesse/précision.
• Maintenir l'écosystème ViT : La compatibilité avec les méthodes de pointe (SSL, TTA, données non-2D) est préservée, ce qui est crucial pour le déploiement pratique et la recherche future.
• Offrir une solution scalable : La méthode fonctionne aussi bien sur des modèles Nano que Large, et s'adapte à divers domaines (vision, séries temporelles, langage).

En conclusion, Jumbo propose une voie simple et efficace pour rendre les Vision Transformers plus rapides et plus puissants sans compromettre leur nature fondamentale, comblant ainsi le fossé entre les modèles "plain" et les architectures optimisées pour l'efficacité.