ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT

Le papier présente ToaSt, un cadre découplé combinant l'élagage structuré des têtes d'attention et la sélection de canaux de tokens pour réduire efficacement les coûts computationnels des Vision Transformers tout en améliorant leur précision.

L'essentiel

🍞 ToaSt : La méthode pour rendre les "cerveaux artificiels" plus légers et plus rapides

Imaginez que vous avez un chef cuisinier génial (c'est le modèle d'intelligence artificielle, ou "Vision Transformer") capable de reconnaître des chats, des voitures ou des visages sur des photos avec une précision incroyable. Le problème ? Ce chef est énorme, il occupe toute la cuisine, il utilise des tonnes d'électricité et il met une éternité à préparer un plat. Il est trop lourd pour être emmené dans un petit camion de livraison (votre téléphone ou une caméra de sécurité).

Les chercheurs ont voulu le rendre plus petit, mais les méthodes habituelles avaient deux gros défauts :

1. La méthode du "ciseau" (pruning) : On coupe des parties du chef. Mais souvent, on le blesse trop, et il faut le faire réapprendre pendant des mois pour qu'il retrouve son niveau. C'est long et cher.
2. La méthode du "résumé" (compression de tokens) : On lui dit de ne regarder que quelques ingrédients au lieu de tous. Mais cela crée un effet domino : si on enlève un ingrédient au début, tout le reste du plat est déséquilibré.

ToaSt (le nom de la nouvelle méthode) est une approche différente. C'est comme si on réorganisait la cuisine en deux étapes distinctes, sans jamais faire réapprendre le chef.

1. La première étape : Le "Coupe-Chou" intelligent (MHSA)

Dans la cuisine du chef, il y a une équipe de 5 assistants (les "têtes" de l'attention) qui regardent tous les détails de l'image en même temps.

• Le problème : Parfois, deux assistants font exactement la même chose. Ils sont redondants.
• La solution ToaSt : Au lieu de licencier un assistant au hasard (ce qui casserait l'équipe), ToaSt regarde les outils de travail de l'équipe. Il remarque que l'assistant A et l'assistant B utilisent les mêmes couteaux pour la même tâche.
• L'action : Il retire les couteaux inutiles de tous les assistants en même temps, de manière synchronisée.
• Le résultat : L'équipe est plus petite, mais elle fonctionne toujours parfaitement ensemble. Pas besoin de réapprendre à travailler !

2. La deuxième étape : Le "Filtre à Bruit" (FFN)

Ensuite, le chef passe à la phase de "cuisson" (le réseau de neurones FFN). C'est là que le gros du travail se fait (plus de 60% de l'énergie est dépensée ici).

• Le problème : Le chef utilise une énorme table pour préparer les ingrédients, mais en réalité, il n'utilise que quelques coins de la table. Le reste est rempli de poussière et de bruit inutile.
• L'observation ToaSt : Les chercheurs ont découvert que plus le chef avance dans sa recette (plus on va profondément dans le réseau), plus il y a de "zones mortes" sur la table. De nombreux ingrédients sont ignorés ou ne servent à rien.
• La solution ToaSt (Sélection de Canaux) : Au lieu de demander au chef de réapprendre, ToaSt utilise un filtre intelligent.
  • Il regarde ce que le chef fait pendant la cuisson.
  • Il repère les ingrédients qui ne servent à rien (le bruit).
  • Il les retire dynamiquement, comme un tamis qui laisse passer l'eau mais garde les pâtes.
• Le résultat : Le chef travaille sur une table plus petite, mais il ne perd aucune information importante. Il devient plus rapide et, étrangement, plus précis car il ne se laisse plus distraire par le bruit.

🚀 Pourquoi c'est révolutionnaire ?

1. Pas de réapprentissage pénible : La plupart des méthodes obligent le chef à réapprendre pendant 300 jours. ToaSt, lui, ne demande que 15 jours de petit ajustement pour les gros modèles. C'est comme passer d'un marathon à une petite promenade.
2. Plus performant qu'avant : En enlevant le "bruit" inutile, le modèle devient souvent meilleur que l'original. Sur l'image de référence (ViT-MAE-Huge), le modèle compressé a obtenu 88,52% de précision (contre 86,88% avant), tout en utilisant 40% moins d'énergie.
3. Ça marche partout : Que ce soit pour reconnaître des chats, des voitures ou pour détecter des objets dans une vidéo, la méthode fonctionne. Même sur des tâches complexes comme la détection d'objets (COCO), le modèle compressé bat le modèle original.

🎯 En résumé

Imaginez que vous avez un camion de déménagement rempli de cartons.

• Les méthodes anciennes : On jette des cartons au hasard, puis on doit tout recharger et réorganiser pendant des heures.
• ToaSt : On regarde les cartons. On vide ceux qui ne contiennent que de la poussière (le bruit). On retire les étagères inutiles du camion (les dimensions inutiles). On garde le camion léger, mais tout le mobilier précieux est toujours là, bien rangé.

ToaSt, c'est la méthode qui permet de faire tenir un cerveau d'intelligence artificielle géant dans votre poche, sans qu'il perde sa mémoire, et même en le rendant plus intelligent ! 🧠✨

Résumé technique

Titre : ToaSt : Sélection de canaux de tokens et élagage structuré pour des ViT efficaces

1. Problématique

Les Transformers de Vision (ViT) ont démontré des performances exceptionnelles dans diverses tâches de vision par ordinateur, surpassant souvent les réseaux de neurones convolutifs (CNN). Cependant, leur déploiement sur des dispositifs aux ressources limitées (mobiles, edge computing) est entravé par des coûts computationnels prohibitifs.

L'analyse de la complexité des ViT révèle deux sources principales de coût :

1. Mécanisme d'attention : Complexité quadratique $O(N^2)$ par rapport à la longueur de la séquence $N$.
2. Réseaux de Neurones Feed-Forward (FFN) : Ces couches représentent environ 61 % des FLOPs totaux (opérations $D^2$), bien que les méthodes d'élagage actuelles se concentrent principalement sur l'attention.

Les solutions existantes souffrent de limitations majeures :

• Élagage de poids structuré : Nécessite un réentraînement (fine-tuning) long et coûteux (souvent >300 époques) et ne cible pas efficacement la redondance des FFN.
• Compression de tokens : Réduit la longueur de séquence $N$ mais ignore la redondance au niveau des canaux (dimension $D$). De plus, les décisions de compression se propagent globalement à travers les couches, créant des dépendances inter-couches qui compliquent l'optimisation.

2. Méthodologie : Le Framework ToaSt

ToaSt propose une approche découplée et indépendante des couches (Layer-Independent) qui applique des stratégies spécialisées distinctes aux deux composants principaux du ViT : les modules d'attention (MHSA) et les réseaux Feed-Forward (FFN).

A. Élagage de poids structuré couplé pour le MHSA (Multi-Head Self-Attention)

• Objectif : Réduire la dimension interne par tête ($d_k$) sans modifier la dimension d'embedding globale ($D$), garantissant la compatibilité avec les connexions résiduelles.
• Mécanisme : L'élagage est effectué de manière synchronisée sur les matrices de poids couplées ($W_Q, W_K, W_V, W_{proj}$). Si une colonne $j$ est supprimée de $W_Q$, la même colonne doit être supprimée de $W_K$ (synchronisation Q-K), et la ligne correspondante de $W_{proj}$ doit être supprimée (synchronisation V-Proj).
• Critère de sélection : Utilisation de la Médiane Géométrique (GM) des poids pré-entraînés. Les dimensions les plus proches du centre de la distribution des poids (les plus redondantes) sont élaguées.
• Stratégie : Élagage uniforme par tête (Head-wise Uniform) avec un taux agressif de 90 % pour toutes les couches sauf la première (qui traite l'interface critique des patchs).

B. Sélection de Canaux de Tokens (TCS) pour les FFN

• Objectif : Réduire la redondance des canaux dans les couches étendues ($D \to 4D$) sans réentraînement.
• Analyse empirique : Les auteurs ont identifié trois signatures de redondance dans les couches profondes :
  1. Fidélité de reconstruction linéaire élevée ($R^2 > 0.9$) : Les canaux sont fortement dépendants les uns des autres.
  2. Effondrement du rang effectif : L'information essentielle réside dans un sous-espace de dimension bien inférieure à $4D$.
  3. Augmentation de la parcimonie (Sparsity) : Une grande partie des neurones sont "morts" (activations proches de zéro) dans les couches profondes.
• Processus (Training-free) :
  • Échantillonnage statistique : Au lieu de calculer l'importance sur tous les tokens (coûteux), un sous-ensemble aléatoire de tokens (2-20 %) est utilisé pour estimer la distribution globale.
  • Métrique unifiée guidée par l'attention : L'importance d'un canal est calculée en combinant l'activation du token CLS (contexte global) et les poids d'attention vers les patches.
  • Élagage asymétrique : Application de ratios conservateurs sur la couche d'expansion (FC1) et agressifs (jusqu'à 90 %) sur la couche de réduction (FC2) dans les couches profondes.

3. Contributions Clés

1. Élagage couplé structuré pour MHSA : Une méthode qui réduit $d_k$ tout en préservant l'intégrité fonctionnelle de l'attention grâce à la synchronisation stricte des indices, permettant une compression indépendante des couches.
2. Analyse de redondance des FFN et TCS : Identification de signatures de redondance (rang, $R^2$, parcimonie) et introduction d'une stratégie de sélection de canaux sans réentraînement qui filtre le bruit redondant.
3. Découplage de l'optimisation : En séparant la compression de l'attention et des FFN, ToaSt évite les paysages d'optimisation complexes et les propagations globales, permettant un élagage agressif avec un coût de récupération minimal.

4. Résultats Expérimentaux

Les évaluations ont été menées sur neuf modèles (DeiT, ViT-MAE, Swin Transformer) sur ImageNet-1K et la tâche de détection d'objets sur COCO.

• Performance sur ImageNet-1K :

  • ViT-MAE-Huge : Atteint 88,52 % de précision Top-1 (gain de +1,64 % par rapport à la base) avec une réduction de 39,4 % des FLOPs.
  • DeiT-Small : Précision de 83,40 % (+3,58 % vs base) avec une accélération de 2,07× sur GPU H100.
  • Efficacité du fine-tuning : Les grands modèles (Huge) nécessitent très peu d'époques de réadaptation (~15 époques) pour récupérer les performances, contrairement aux modèles plus petits qui en nécessitent beaucoup plus.

• Tâches en aval (COCO Object Detection) :

  • Avec un backbone Swin-Base compressé, le modèle atteint 52,2 mAP (boîtes) contre 51,9 mAP pour la base non compressée, démontrant que la compression élimine du bruit plutôt que des caractéristiques discriminatives.

• Comparaison avec l'état de l'art :

  • ToaSt surpasse systématiquement les méthodes de compression de tokens (ToMe, DiffRate) en termes de compromis précision/efficacité, offrant une densité représentationnelle supérieure à budget FLOPs égal.

5. Signification et Impact

ToaSt représente une avancée significative dans la compression des modèles de vision.

• Efficacité sans précédent : Il démontre qu'il est possible de réduire massivement la complexité computationnelle (FLOPs) tout en améliorant la précision du modèle, en agissant comme un régularisateur implicite qui élimine le bruit.
• Adaptabilité : La méthode fonctionne efficacement sur une large gamme d'architectures (DeiT, MAE, Swin) et de tailles de modèles.
• Déploiement pratique : L'absence de réentraînement coûteux pour la partie FFN et la structure régulière des poids élagués (compatible avec les GPU standards sans bibliothèques de matrices creuses spécialisées) facilitent grandement le déploiement sur le matériel réel.

En conclusion, ToaSt propose un paradigme de compression "découplé" qui résout les goulots d'étranglement computationnels des ViT tout en préservant, voire en améliorant, leurs capacités de représentation.