Adaptive MLP Pruning for Large Vision Transformers

Cet article propose une méthode d'élagage adaptif des perceptrons multicouches (AMP) qui réduit d'environ 40 % les paramètres et les FLOPs des grands transformateurs de vision sans perte significative de performance, en utilisant un critère d'entropie informationnelle sans étiquette pour évaluer l'importance des neurones et un algorithme de recherche binaire pour déterminer automatiquement le taux de compression.

L'essentiel

🖼️ Le Dilemme des Géants : Comment rendre les "Géants" de l'IA plus agiles ?

Imaginez que vous avez un chef cuisinier légendaire (c'est le modèle d'IA, comme un "Vision Transformer") capable de reconnaître n'importe quel plat ou image au monde avec une précision incroyable. C'est un génie, mais il y a un problème : sa cuisine est gigantesque, remplie de milliers d'ustensiles inutiles, et il faut des heures pour qu'il prépare un simple plat. C'est ce qu'on appelle un modèle "lourd" : il consomme énormément d'énergie et de temps.

Les chercheurs de l'Université Centrale du Sud (Chengchao Shen) se sont demandé : "Comment pouvons-nous garder ce chef génial, mais le rendre plus rapide et plus léger, sans qu'il perde son talent ?"

Leur réponse s'appelle AMP (Adaptive MLP Pruning), ou en français : Élagage Adaptatif des Couches de Neurones.

Voici comment ça marche, étape par étape, avec des analogies simples :

1. Le Problème : Trop de "Cerveaux" qui ne servent à rien

Dans ces modèles d'IA, il y a une partie spécifique appelée MLP (Perceptron Multicouche). C'est comme le "cerveau" qui réfléchit et analyse les images.

• L'analogie : Imaginez que ce cerveau est composé de 100 000 employés. En réalité, les chercheurs ont découvert que 80 % de ces employés ne font rien d'utile ! Ils sont là, assis, mais ne contribuent pas vraiment à la décision finale. C'est du gaspillage.

2. La Solution : Qui est le meilleur employé ? (L'Évaluation)

Pour savoir quels employés (neurones) on peut virer, il faut les tester.

• L'ancienne méthode (Le test à l'aveugle) : Les méthodes précédentes regardaient seulement si l'employé avait raison sur la réponse exacte (par exemple : "Est-ce que c'est un chat ?"). Si l'employé disait "Non, c'est un chien", on le jugeait mauvais, même s'il avait raison sur le fait que ce n'est pas un chat. C'était un test imparfait.
• La méthode AMP (La carte complète) : Les chercheurs ont inventé un nouveau test basé sur l'"Entropie de l'Information".
  • L'analogie : Au lieu de demander "Est-ce un chat ?", ils demandent à l'employé de décrire toutes les possibilités ("C'est peut-être un chat, un chien, ou un lapin").
  • En regardant toute la carte des possibilités, on voit mieux qui comprend vraiment la situation. Cela permet d'identifier avec une précision chirurgicale les employés qui sont vraiment inutiles, sans se tromper. De plus, ce test fonctionne même si on ne connaît pas la "recette secrète" (les étiquettes) du modèle original.

3. La Coupe : La "Chasse aux Neurons" Intelligente

Une fois qu'on a classé les employés du plus utile au moins utile, il faut décider combien en licencier.

• L'ancienne méthode : On disait "On vire 40 % des employés, tout le monde, partout". C'est comme couper 40 % des branches d'un arbre au hasard. Ça peut tuer l'arbre si certaines branches sont vitales.
• La méthode AMP (Recherche Binaire) : Ils utilisent un algorithme de recherche binaire (comme un jeu de "Plus grand, plus petit").
  • L'analogie : Imaginez que vous cherchez le point idéal pour couper une branche. Vous essayez de couper un peu, puis vous vérifiez si l'arbre tient toujours debout. Si oui, vous coupez encore un peu. Si l'arbre tremble trop, vous remettez un peu de bois.
  • Le système s'adapte : il coupe beaucoup là où il y a du gaspillage, et peu là où c'est critique. Il ne fixe pas de pourcentage à l'avance ; il trouve le point de rupture parfait pour chaque partie du modèle.

4. La Rééducation : Le Mentorat (Distillation de Connaissances)

Après avoir viré les employés inutiles, le chef cuisinier (le modèle réduit) est un peu perdu. Il a moins de monde pour l'aider.

• L'analogie : Le "Grand Chef" (le modèle original, énorme) reste là pour mentoriser le "Petit Chef" (le modèle réduit).
• Le Grand Chef dit au Petit Chef : "Regarde, quand je vois cette image, je pense à ça. Toi, tu devrais penser pareil."
• Grâce à ce coaching, le Petit Chef récupère presque instantanément toutes les compétences du Grand Chef, mais en étant beaucoup plus rapide et léger.

🏆 Les Résultats Magiques

Grâce à cette méthode, les chercheurs ont obtenu des résultats impressionnants sur des modèles très célèbres (comme CLIP et DINOv2) :

• Réduction de poids : Ils ont supprimé environ 40 % des paramètres (les "employés" inutiles).
• Vitesse : Le modèle est 1,5 fois plus rapide.
• Performance : Le plus étonnant ? Le modèle réduit fait aussi bien (voire un tout petit peu mieux) que le modèle original !
• Sans étiquettes : Ils ont pu faire cela même sur des modèles dont on ne connaissait pas toutes les données d'entraînement (comme DINOv2), grâce à leur test "sans étiquette".

En résumé

C'est comme si vous preniez une bibliothèque de 10 000 livres, vous lisiez chaque page pour voir quelles histoires sont vraiment importantes, vous jetez les 4 000 livres inutiles, puis vous demandez au bibliothécaire original de résumer les 6 000 livres restants pour que le nouveau bibliothécaire (plus petit) puisse répondre aux questions aussi bien que l'ancien, mais en courant beaucoup plus vite.

C'est une avancée majeure pour rendre l'intelligence artificielle plus accessible, moins coûteuse en énergie et plus rapide à utiliser sur nos téléphones et ordinateurs.

Résumé technique

1. Problématique

Les grands Transformers Visuels (Large Vision Transformers), tels que CLIP et DINOv2, démontrent des performances exceptionnelles grâce à une mise à l'échelle accrue de leur capacité. Cependant, cette performance s'accompagne d'une demande computationnelle et mémoire excessive, rendant leur déploiement coûteux et difficile.

L'analyse de l'architecture de ces modèles révèle que les modules Multilayer Perceptron (MLP) constituent la majorité des paramètres (par exemple, 81,1 % dans EVA-CLIP-E). Par conséquent, l'élagage (pruning) ciblé de ces modules offre une opportunité majeure de compression.

Les méthodes d'élagage existantes, notamment celles basées sur l'expansion de Taylor, souffrent de deux limitations principales :

1. Critère d'évaluation imparfait : Elles utilisent généralement une perte de cross-entropy one-hot pour évaluer l'importance des neurones. Cette approche ignore les probabilités de prédiction pour les autres classes, ce qui réduit la fidélité du score d'importance.
2. Dépendance aux étiquettes et aux modules externes : Ces méthodes nécessitent souvent les étiquettes du jeu de données ou des modules spécifiques (comme la tête DINO ou l'encodeur texte de CLIP) qui ne sont pas toujours disponibles publiquement (ex: DINOv2).
3. Rigidité du taux de compression : La plupart des méthodes imposent un taux d'élagage prédéfini, sans s'adapter à la redondance spécifique de chaque module MLP.

2. Méthodologie : AMP (Adaptive MLP Pruning)

Les auteurs proposent une méthode d'Élagage Adaptatif des MLP (AMP) composée de trois étapes clés :

A. Évaluation de l'importance des neurones via l'Entropie d'Information

Au lieu de la cross-entropy one-hot, les auteurs introduisent un critère sans étiquette (label-free) basé sur l'entropie d'information.

• Principe : Pour chaque neurone caché du MLP, l'importance est évaluée en mesurant la variation de l'entropie des prédictions du modèle lorsque ce neurone est supprimé.
• Avantage : Cette méthode exploite la distribution complète des probabilités de prédiction (toutes les classes), offrant une estimation plus précise de l'importance.
• Indépendance : Elle ne dépend pas de la fonction de perte originale ni des modules de tête (ex: DINO head), permettant l'élagage de modèles dont les poids complets ne sont pas publiés. La similarité inter-instance est utilisée pour reconstruire les probabilités de prédiction sans étiquettes.

B. Élagage Adaptatif par Recherche Binaire

Une fois les neurones classés par importance, la méthode ne supprime pas un pourcentage fixe.

• Algorithme : Une recherche binaire est appliquée sur chaque module MLP pour déterminer le nombre optimal de neurones à conserver.
• Critère d'arrêt : Le processus itère en ajustant la taille du MLP jusqu'à ce que la variation de l'entropie d'information du modèle élagué par rapport au modèle original dépasse un seuil tolérable ($\Delta E$).
• Résultat : Cela permet d'adapter le taux de compression à la redondance spécifique de chaque couche, évitant un taux de compression prédéfini arbitraire.

C. Récupération des Performances par Distillation de Connaissances

Pour compenser la perte de capacité due à l'élagage :

• Le modèle original agit comme un enseignant (teacher) et le modèle élagué comme un élève (student).
• Une distillation de connaissances est effectuée en minimisant l'erreur quadratique moyenne (MSE) entre les sorties des blocs finaux (jetons de classe et jetons de patch) des deux modèles.
• Grâce à la similarité structurelle (les dimensions de sortie restent identiques car seuls les neurones internes sont supprimés), aucune couche d'alignement supplémentaire n'est nécessaire.

3. Contributions Clés

1. Critère d'Entropie d'Information : Introduction d'un critère d'évaluation d'importance précis et sans étiquettes, permettant la compression de modèles dont le code source ou les poids complets ne sont pas ouverts.
2. Élagage Adaptatif : Développement d'une stratégie de recherche binaire qui élimine les neurones redondants de manière dynamique, évitant les taux de compression prédéfinis.
3. Accélération Quasi Sans Perte : Réduction d'environ 40 % des paramètres et des FLOPs avec une récupération quasi parfaite des performances après distillation.
4. Supériorité sans Fine-tuning : Même sans fine-tuning (distillation), la méthode surpasse largement les autres méthodes d'élagage.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles de pointe (OpenCLIP, EVA-CLIP, DINOv2) sur plusieurs benchmarks (ImageNet-1K, ObjectNet, Flickr30K, COCO).

• Compression : Réduction d'environ 40 % des paramètres et des opérations (FLOPs).
• Accélération : Gain d'environ 1,5x dans le débit d'inférence (images/seconde).
• Performance (Zero-shot) :
  • Après distillation, les modèles élagués récupèrent presque totalement les performances des modèles originaux. Dans certains cas (ex: OpenCLIP-g, EVA-CLIP-E), ils surpassent légèrement les modèles originaux.
  • Sur la tâche de classification zero-shot, le modèle élagué OpenCLIP-g atteint 73,1 % de précision (vs 73,0 % original) avec 40 % de paramètres en moins.
• Comparaison : La méthode AMP surpasse significativement les méthodes de référence (Random pruning, $\ell_2$ norm, Taylor pruning, SAViT, NViT). Par exemple, sans fine-tuning, AMP obtient 53,8 % de précision sur OpenCLIP-g contre 11,1 % pour NViT.
• Validité sur Transformers Purs : La méthode fonctionne également sur DINOv2-g (un modèle sans tête de classification explicite), confirmant l'efficacité du critère sans étiquette.

5. Signification et Impact

Ce travail propose une avancée significative pour le déploiement efficace des grands modèles de vision.

• Accessibilité : En rendant possible l'élagage de modèles dont les composants de prédiction ne sont pas publics (comme DINOv2), la méthode élargit l'accès à la compression de modèles propriétaires ou partiellement ouverts.
• Efficacité : Elle démontre qu'une grande partie des paramètres des MLP dans les grands Transformers est redondante et peut être éliminée dynamiquement sans sacrifier la précision.
• Futur : Les auteurs prévoient d'étendre cette approche adaptative aux modules d'attention multi-têtes et aux grands modèles de langage (LLM).

En résumé, l'AMP offre une solution robuste, adaptable et hautement efficace pour réduire l'empreinte des grands Transformers visuels tout en maintenant, voire en améliorant, leurs capacités de généralisation.