HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

Ce papier présente HiAP, un cadre d'élagage stochastique multi-granulaire qui optimise de manière continue et en une seule phase les Transformers de vision en éliminant simultanément des têtes d'attention, des blocs FFN et des dimensions internes, permettant ainsi de réduire efficacement les coûts computationnels et mémoire pour le déploiement sur appareils embarqués sans nécessiter de seuils manuels ou de pipelines complexes.

L'essentiel

Imaginez que vous avez un énorme camion de déménagement (le Vision Transformer) qui doit transporter des meubles précieux (les images) à travers la ville. Ce camion est très puissant, mais il est aussi énorme, gourmand en essence et difficile à garer dans les petites ruelles de nos téléphones ou de nos montres connectées (les appareils mobiles).

L'objectif de cette recherche est de transformer ce camion géant en un véhicule agile et économe, sans perdre de meubles en cours de route.

Voici comment les auteurs, avec leur méthode appelée HiAP, y parviennent, expliquée simplement :

1. Le Problème : Un camion trop lourd

Les camions actuels (les modèles d'IA) sont construits avec une structure rigide. Pour les rendre plus petits, les méthodes précédentes faisaient deux choses :

• Soit elles enlevaient des roues individuelles (ce qui rendait le camion instable et difficile à réparer).
• Soit elles enlevaient des étages entiers, mais de manière très brutale, en demandant à un expert humain de dire : "Enlève cet étage, garde celui-là". C'était long, compliqué et nécessitait souvent de repartir de zéro pour réparer les dégâts.

2. La Solution HiAP : Le Camion "Chaméléon"

HiAP est comme un camion de déménagement magique qui peut se remodeler tout seul pendant qu'il roule. Au lieu d'avoir un chauffeur humain qui décide quoi jeter, le camion possède une intelligence intégrée qui apprend à devenir plus petit et plus efficace tout en apprenant sa tâche.

Le secret réside dans deux niveaux de contrôle, comme des interrupteurs à deux tailles :

A. Les Gros Interrupteurs (Macro) : "On coupe un étage entier !"

Imaginez que votre camion a plusieurs étages (des blocs de traitement) et plusieurs équipes de déménageurs (les "têtes d'attention").

• Parfois, un étage entier est inutile. HiAP peut dire : "On n'a pas besoin de cet étage, on le retire complètement !"
• Parfois, une équipe de déménageurs est de trop. HiAP peut dire : "On envoie cette équipe chez elle, on ne garde que les autres."
• L'avantage : Cela réduit énormément le poids du camion et la consommation d'essence (mémoire), car on ne charge même plus les étages inutiles.

B. Les Petits Interrupteurs (Micro) : "On réduit la taille des outils"

Si on ne peut pas retirer un étage entier, HiAP regarde à l'intérieur.

• Au lieu d'avoir 100 outils dans la boîte à outils d'un déménageur, il décide : "Tu n'as besoin que de 50 outils pour ce travail."
• Il retire les outils inutiles, mais garde le déménageur en place.
• L'avantage : Cela rend le camion plus rapide à manœuvrer (calculs mathématiques) sans changer sa structure de base.

3. La Magie : Apprendre en une seule fois

La plupart des méthodes précédentes ressemblaient à un processus en deux étapes :

1. Étape 1 : Un expert essaie de deviner quoi enlever (c'est long et difficile).
2. Étape 2 : On doit réapprendre à conduire le camion avec les nouveaux meubles (c'est encore plus long).

HiAP change la donne : C'est un apprentissage en une seule étape.

• Le camion commence avec tous ses étages et tous ses outils.
• Pendant l'entraînement, il utilise une technique spéciale (des "portes aléatoires" qui deviennent de plus en plus fermées) pour tester ce qui est vraiment nécessaire.
• Au fur et à mesure qu'il apprend, il se rend compte : "Tiens, je n'ai pas besoin de cet étage" ou "Je n'ai besoin que de la moitié de ces outils".
• À la fin de l'entraînement, le camion est déjà petit, léger et prêt à rouler. Aucune réparation n'est nécessaire.

4. Pourquoi c'est génial ?

• Pas de chef d'orchestre : Le modèle décide tout seul de quoi garder et de quoi jeter. Pas besoin d'experts humains pour régler des boutons compliqués.
• Économie double : Il économise à la fois l'espace de stockage (en enlevant les étages) et l'énergie de calcul (en réduisant les outils). C'est comme si votre voiture consommait moins d'essence ET prenait moins de place dans le garage.
• Résultat : Sur des tests réels (comme reconnaître des images), HiAP a réussi à créer des camions beaucoup plus petits qui font presque aussi bien que les géants originaux, mais qui sont beaucoup plus rapides sur les petits appareils.

En résumé : HiAP est un architecte intelligent qui construit une maison (le modèle) en apprenant à la fois à enlever les pièces inutiles et à réduire la taille des meubles, le tout pendant la construction, pour obtenir une maison parfaite, petite et efficace, sans jamais avoir besoin de la démolir pour la reconstruire.

Résumé technique

1. Problématique

Les Vision Transformers (ViT) dominent actuellement le domaine de la vision par ordinateur, mais leur déploiement sur des dispositifs périphériques (edge devices) est sévèrement limité par leurs coûts computationnels élevés et leur consommation de bande passante mémoire.

Les méthodes d'élagage (pruning) existantes souffrent de deux limitations majeures :

• Granularité unique : La plupart des méthodes opèrent soit au niveau macro (élagage de têtes d'attention ou de blocs entiers), soit au niveau micro (réduction des dimensions internes ou des neurones). L'élagage micro préserve la profondeur et le nombre de têtes, ce qui ne résout pas le goulot d'étranglement lié à la bande passante mémoire (chargement de matrices complètes). À l'inverse, l'élagage macro réduit la mémoire mais risque de dégrader fortement la capacité de représentation du modèle.
• Processus complexe et manuel : Les approches actuelles reposent souvent sur des pipelines multi-étapes, des heuristiques manuelles pour déterminer l'importance des composants, et un seuillage post-hoc (après entraînement) pour satisfaire les budgets de sparsité. Cela nécessite une intervention experte et des phases de fine-tuning coûteuses.

2. Méthodologie : HiAP (Hierarchical Auto-Pruning)

HiAP propose un cadre d'apprentissage end-to-end et en une seule phase qui découvre automatiquement la sous-réseau optimale sans seuillage manuel ni cibles de sparsité prédéfinies.

A. Mécanisme de Portes Hiérarchique (Gating)

HiAP introduit des portes stochastiques basées sur la distribution Gumbel-Sigmoid à deux niveaux de granularité simultanés :

1. Niveau Macro : Des portes contrôlent la rétention de structures entières (têtes d'attention complètes et blocs FFN). Si une porte est désactivée, le bloc ou la tête est entièrement contourné, éliminant ainsi le surcoût mémoire de chargement.
2. Niveau Micro : Au sein des structures macro actives, des portes sélectionnent dynamiquement des dimensions spécifiques au sein des têtes d'attention et des neurones individuels dans les couches Feed-Forward (FFN). Cela permet d'ajuster la capacité de calcul (FLOPs) sans supprimer la structure entière.

B. Modélisation du Coût Différentiable

Le modèle intègre une formulation analytique exacte des opérations Multiply-Accumulate (MACs) directement dans la fonction de perte.

• Le coût est décomposé en une surcharge macro (calcul des projections Q, K, V et de la carte d'attention) et des coûts micro (calcul des sorties de dimensions ou de neurones survivants).
• Cette décomposition linéaire permet d'attribuer des pénalités précises aux structures vides, incitant le modèle à supprimer les têtes entières plutôt que de garder des têtes avec des dimensions inefficaces.

C. Optimisation et Contraintes de Faisabilité

• Fonction de Perte : Elle combine la perte de tâche (avec distillation de connaissances), une régularisation consciente du budget (pénalisant les MACs attendus) et des pénalités de faisabilité structurelle.
• Prévention de l'effondrement : Des contraintes explicites (via des fonctions ReLU) garantissent un nombre minimum de têtes et de dimensions par couche, empêchant le réseau de supprimer des couches entières prématurément pour minimiser artificiellement le coût.
• Recherche en une phase : Contrairement aux méthodes NAS (Neural Architecture Search) classiques qui nécessitent une phase de recherche suivie d'un fine-tuning, HiAP entraîne le réseau dense et les portes simultanément. La température $\tau$ de la distribution Gumbel-Sigmoid est progressivement réduite (recuit), transformant les portes continues en décisions binaires déterministes à la fin de l'entraînement.

3. Contributions Clés

1. Unification Multi-Niveau : HiAP est le premier cadre à unifier l'élagage macro (têtes/blocs) et micro (dimensions/neurones) dans un seul système de portes différentiable, adressant simultanément les goulots d'étranglement de la mémoire et du calcul.
2. Découverte Autonome : Élimination totale des heuristiques manuelles, des métriques de classement proxy et des phases de fine-tuning secondaires. Le modèle apprend sa propre architecture optimale sous contrainte de budget.
3. Sous-réseau Physiquement Extractible : À la fin de l'entraînement, le modèle produit un sous-réseau dense et compressé physiquement (les matrices sont tronquées), prêt pour l'inférence sur du matériel standard sans moteurs de convolution sparse.

4. Résultats Expérimentaux

Les expériences ont été menées sur CIFAR-10 (ViT-Tiny) et ImageNet-1K (DeiT-Small).

• Performance sur ImageNet : HiAP a compressé DeiT-Small de 4,6G à 3,1G MACs (réduction de ~33%) tout en maintenant une précision Top-1 de 79,10 %. Ce résultat est compétitif par rapport aux méthodes de pointe (SOTA) comme GOHSP ou ViT-Slim, qui nécessitent des pipelines complexes.
• Efficacité sur CIFAR-10 : HiAP surpasse les heuristiques d'élagage standard (comme le tri par norme L1 ou les ratios uniformes), offrant une amélioration de précision de +0,93 % à budget égal.
• Gain de Latence : Sur un GPU, le modèle élagué de 33 % a montré une réduction de la latence d'inférence de 5,57 ms à 3,86 ms, soit un gain de vitesse d'environ 1,44x, confirmant que les gains théoriques se traduisent par des accélérations réelles sur du matériel standard.
• Dynamique d'Apprentissage : L'analyse montre que HiAP élimine d'abord les structures macro (têtes entières, blocs FFN redondants) pour réduire la surcharge mémoire, puis affine la sparsité au niveau micro pour optimiser les FLOPs résiduels.

5. Signification et Impact

Ce travail représente une avancée significative pour le déploiement des Transformers sur des dispositifs contraints :

• Simplification du Pipeline : Il transforme un problème d'ingénierie complexe (multi-étapes, réglage manuel) en un problème d'apprentissage automatique unifié.
• Optimisation Holistique : En traitant conjointement la bande passante mémoire (via l'élagage macro) et le calcul (via l'élagage micro), HiAP offre une frontière de Pareto précision-efficacité supérieure.
• Déploiement Pratique : La capacité à générer un sous-réseau dense et physiquement compressé sans fine-tuning postérieur facilite grandement l'intégration dans les systèmes embarqués et les applications temps réel.

En résumé, HiAP démontre qu'il est possible de découvrir automatiquement des architectures de Vision Transformers hautement efficaces, équilibrant parfaitement la réduction de la mémoire et du calcul, tout en garantissant la stabilité de l'entraînement et la performance finale.