🤖 AI

Learning to Predict Middle-Layer Attention in MLLMs for Visual Token Prunin

Cet article propose la Prédiction d'Attention en Couche Intermédiaire (MAP), une méthode qui identifie dynamiquement les couches d'attention optimales spécifiques à chaque échantillon et les distille dans un prédicteur léger pour élaguer les jetons visuels avant le traitement par le modèle de langage, atteignant une accélération de bout en bout de 3,09x tout en conservant 97,5 % des performances sur LLaVA-NeXT-7B.

Auteurs originaux : Yuyao Sun, Tao Deng, Shuang Li, Deqing Wang, Hao Geng, Minjun Yu

Publié 2026-08-10

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Yuyao Sun, Tao Deng, Shuang Li, Deqing Wang, Hao Geng, Minjun Yu

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'apprendre à un robot super intelligent à voir et à parler du monde en même temps. C'est le monde des Modèles de Langage de Grande Taille Multimodaux (MLLM). Voyez ces modèles comme des détectives brillants capables de regarder une image et de répondre à des questions à son sujet, comme « De quelle couleur est le collier du chien ? » ou « Pourquoi la personne court-elle ? ». Pour ce faire, le robot ne se contente pas de « regarder » l'image ; il la décompose en des milliers de minuscules pièces de puzzle numériques appelées tokens visuels. Chaque token est un petit fragment d'information sur une partie de l'image.

Le problème est que ces robots sont incroyablement affamés. Pour obtenir une vraiment bonne vue d'une image, ils peuvent avoir besoin de traiter des milliers de ces tokens. C'est comme essayer de lire une encyclopédie entière pour trouver la réponse à une question simple sur le chapeau d'un chat. Cela prend beaucoup de temps et d'énergie (puissance de calcul), ce qui rend le robot lent et coûteux à exploiter. Les scientifiques ont essayé de trouver un moyen de rendre ces robots plus rapides en jetant les pièces de puzzle inutiles avant même qu'ils ne commencent à les lire. C'est ce qu'on appelle l'élagage des tokens visuels (visual token pruning). L'astuce consiste à savoir quels morceaux garder et lesquels jeter. Si vous jetez le mauvais morceau, le robot est confus. Si vous en gardez trop, il reste lent. La grande question a été : comment savoir exactement quels morceaux de l'image sont les plus importants pour la question spécifique posée ?

Le Problème : Une taille unique ne convient pas à tous

Les chercheurs derrière ce papier, dirigés par Yuyao Sun et Tao Deng, ont découvert une faille délicate dans la manière dont nous essayons actuellement d'aider ces robots à décider quoi garder.

Auparavant, les scientifiques pensaient que si l'on observait l'« attention » du robot (une façon élégante de dire « ce sur quoi le robot se concentre ») à partir d'une couche spécifique, située au milieu, on pourrait trouver les pièces d'image les plus importantes. C'était comme dire : « Hé robot, regarde simplement ce sur quoi tu réfléchissais à la moitié de ton processus de pensée, et garde ces pièces d'image ».

Mais les auteurs ont découvert que cette approche « taille unique » est défaillante. Imaginez que vous regardiez une photo de plage.

Si quelqu'un demande : « Quelle est la couleur du sable ? », le cerveau du robot s'illumine sur le sable au milieu de son processus de réflexion.
Si quelqu'un demande : « Quelle est la couleur du ciel ? », le cerveau du robot s'illumine sur le ciel, mais peut-être à un stade de sa réflexion différent.

Les chercheurs ont montré que la « meilleure » couche du cerveau à vérifier dépend entièrement de la question spécifique. Parfois, la réponse se trouve dans la couche 10, parfois dans la couche 20. Utiliser une couche fixe revient à essayer d'utiliser une seule clé pour toutes les portes d'un bâtiment ; cela fonctionne pour certaines, mais échoue pour d'autres.

De plus, il y avait un second problème, plus important encore. Pour découvrir quelle couche était la « meilleure » pour une question spécifique, le robot devait d'abord traiter l'image entière à travers toutes ces couches. C'est comme essayer de trouver le meilleur itinéraire pour aller à l'école en faisant d'abord tout le trajet, juste pour voir quel virage était le meilleur. Au moment où vous déterminez quels morceaux garder, vous avez déjà gaspillé tout le temps et l'énergie que vous essayiez de gagner !

La Solution : Le Détective du « Contraste de Question »

Pour résoudre cela, l'équipe a proposé une nouvelle méthode appelée MAP (Middle-layer Attention Prediction - Prédiction de l'Attention de la Couche Médiane). Au lieu de faire subir au robot le travail de force pendant la conversation réelle, ils ont appris à un petit assistant, super rapide, à deviner ce sur quoi le robot se serait concentré.

Voici comment ils ont entraîné cet assistant :

L'astuce du « Contraste » (QCTS) : Ils ont trouvé une manière ingénieuse de choisir le bon « enseignant » pour chaque question. Ils ont pris une image et ont posé deux questions au robot :
- Question A : La vraie question (ex: « Quelle est la couleur du bateau ? »).
- Question B : Une question banale et générique (ex: « Qu'y a-t-il dans cette image ? »).
Ils ont ensuite comparé l'attention du robot pour les deux questions. Les couches où l'attention du robot a le plus changé entre les deux questions sont celles qui se sont réellement souciées des détails spécifiques du bateau. Cette méthode, appelée Question Contrastive Teacher Selection (QCTS), agit comme un projecteur, identifiant instantanément quelle partie du cerveau du robot est la plus excitée par la question spécifique.
Le Prédicteur Léger : Une fois qu'ils ont su quelle couche était la « star » pour une question donnée, ils n'ont pas gardé cette couche lourde. À la place, ils ont appris à un prédicteur minuscule et léger à imiter le comportement de cette couche star. Ce prédicteur est si petit et rapide qu'il peut regarder l'image et la question avant même que le gros robot ne commence à réfléchir, et dire immédiatement : « Gardez ces 5 % de pièces d'image, jetez le reste ».

Les Résultats : Rapide, Léger et Intelligent

L'équipe a testé ce nouveau système sur plusieurs cerveaux de robots différents (MLLM) et une grande variété de questions complexes. Les résultats sont impressionnants.

Vitesse : Sur l'un de leurs modèles de test (LLaVA-NeXT-7B), MAP a rendu le robot 3,09 fois plus rapide du début à la fin. C'est comme transformer une marche de 10 minutes en un jogging de 3 minutes.
Efficacité : Ils ont réussi à jeter 94,4 % des tokens visuels (en ne gardant que 5,56 % d'entre eux) et le robot a tout de même obtenu 97,5 % des bonnes réponses qu'il aurait obtenues avec l'image complète.
Sélection Intelligente : Contrairement à d'autres méthodes qui se contentent de deviner ou d'utiliser une règle fixe, la règle de « diversité » de MAP a permis de s'assurer que, même en gardant très peu de tokens, elle ne choisissait pas seulement des pièces qui se ressemblent. Elle a choisi des pièces différentes les unes des autres mais toujours pertinentes pour la question, garantissant que le robot ne manque aucun détail crucial.

Pourquoi cela compte

Ce papier ne se contente pas de suggérer une idée cool ; il fournit un outil opérationnel qui résout un véritable goulot d'étranglement. En prouvant que la « meilleure » couche change pour chaque question et que nous pouvons la prédire sans effectuer le travail lourd au préalable, MAP permet à ces puissants modèles d'IA de fonctionner beaucoup plus rapidement sur des ordinateurs standards. C'est un peu comme donner à un bibliothécaire une fiche d'index magique qui lui indique exactement quel livre tirer de l'étagère avant même qu'il ne pénètre dans la bibliothèque, lui évitant ainsi de devoir scanner chaque livre sur le mur.

Les auteurs montrent qu'avec cette méthode, nous pouvons avoir le beurre et l'argent du beurre : nous pouvons maintenir une intelligence élevée du robot tout en le rendant nettement plus rapide et moins coûteux à exploiter. Cela suggère que l'avenir de l'IA ne réside pas seulement dans la création de cerveaux plus gros et plus lours, mais dans l'apprentissage de la manière d'être plus intelligents dans la façon dont ils prêtent attention.

Résumé Technique : Apprendre à prédire l'attention des couches intermédiaires dans les MLLM pour l'élagage des jetons visuels

Énoncé du Problème

Les modèles de langage multimodaux (MLLM) ont démontré de fortes capacités dans les tâches de vision-langage mais souffrent de coûts de calcul élevés en raison du traitement de nombreux jetons visuels (visual tokens). Bien que l'élagage des jetons visuels offre une solution pour améliorer l'efficacité de l'inférence, les méthodes existantes font face à deux limitations critiques :

Couches Fixes Sous-optimales : Des études récentes suggèrent que l'attention texte-vers-vision des couches intermédiaires du modèle de langage guide efficacement l'élagage des jetons. Cependant, l'analyse révèle que la couche spécifique fournissant l'attention la plus réactive à la question varie considérablement d'un échantillon à l'autre. Une stratégie à couche fixe est donc sous-optimale.
Dilemme de l'Efficacité d'Inférence : L'extraction de l'attention d'une couche intermédiaire appropriée nécessite de faire passer de nombreux jetons visuels à travers plusieurs couches du modèle de langage avant que l'élagage ne puisse avoir lieu. Cela engendre un calcul substantiel avant toute suppression de jetons, annulant les gains potentiels d'efficacité. De plus, la matérialisation explicite des cartes d'attention pendant l'inférence peut limiter la compatibilité avec les noyaux d'attention optimisés (ex: FlashAttention).

Méthodologie : Prédiction de l'Attention des Couches Intermédiaires (MAP)

Pour répondre à ces problèmes, les auteurs proposent MAP, un cadre qui distille l'attention des couches intermédiaires spécifique à l'échantillon dans un prédicteur léger opérant avant la première couche du modèle de langage. MAP se compose de deux phases :

1. Entraînement Hors Ligne (Offline Training)

Sélection de l'Enseignant par Contraste de Question (QCTS) : Au lieu d'utiliser une couche fixe, MAP sélectionne dynamiquement une couche "enseignante" pour chaque échantillon d'entraînement. Il compare la distribution de l'attention texte-vers-vision générée par la question d'entrée contre une question de référence générique ("Que voit-on sur cette image ?") sur la même image. La couche présentant la plus grande divergence de Jensen–Shannon entre ces deux distributions est sélectionnée comme enseignant, car elle capture le mieux les régions pertinentes pour la question spécifique.
Distillation de l'Attention Cross-Modale : Un prédicteur léger est entraîné pour reproduire la distribution d'attention de l'enseignant sélectionné. Le prédicteur prend les représentations d'entrée multimodales (jetons visuels du codeur de vision et jetons textuels de la couche d'embedding) disponibles avant le traitement par le LLM. Il utilise des MLP spécifiques aux modalités et des couches de projection pour estimer les scores d'attention, minimisant la divergence de Kullback-Leibler (KL) entre la distribution prédite et la distribution de l'enseignant. Le MLLM original reste gelé durant ce processus.

2. Élagage des Jetons au Moment de l'Inférence

Prédiction et Sélection : Pendant l'inférence, le prédicteur estime la distribution de l'attention texte-vers-vision directement à partir des caractéristiques multimodales. Ces scores servent d'importance pour les jetons visuels.
Élagage Sensible à la Diversité : Pour garantir que les jetons conservés soient non seulement pertinents mais aussi complémentaires, MAP combine les scores d'importance prédits avec un critère de diversité. Il sélectionne les jetons de manière gloutonne, en commençant par les $K_0$ les plus importants, et ajoute de manière itérative des candidats qui maximisent le résidu orthogonal par rapport au sous-espace actuel des plongements (embeddings) sélectionnés.
Exécution : L'élagage se produit avant la première couche du modèle de langage. Les jetons sélectionnés sont restaurés dans leur ordre d'origine et combinés avec les jetons de la question, permettant au traitement ultérieur d'utiliser uniquement l'ensemble compact de jetons. Cette approche évite de dépendre des cartes d'attention du LLM, maintenant la compatibilité avec les noyaux optimisés.

Principales Contributions

Analyse de la Variabilité des Couches : Les auteurs démontrent que bien que l'attention texte-vers-vision des couches intermédiaires soit efficace pour l'élagage, la couche optimale dépend de l'échantillon, rendant les stratégies à couche fixe sous-optimales.
Cadre MAP : Ils proposent une méthode qui distille l'attention spécifique à l'échantillon dans un prédicteur, permettant un élagage efficace avant le traitement par le LLM sans subir le coût de la traversée de multiples couches de langage.
Mécanisme QCTS : Ils introduisent la Sélection de l'Enseignant par Contraste de Question (QCTS), qui identifie la couche la plus réactive à la question pour chaque échantillon en contrastant les distributions d'attention avec une question de référence, fournissant ainsi une supervision sur mesure pour le prédicteur.
Validation Empirique : Des expériences approfondies sur trois backbones de MLLM (LLaVA-1.5-7B, LLaVA-NeXT-7B, Qwen2.5-VL-7B) et dix benchmarks montrent que MAP atteint un compromis précision-efficacité supérieur aux méthodes d'élagage existantes.

Résultats Expérimentaux

L'article rapporte des résultats sur LLaVA-NeXT-7B et LLaVA-1.5-7B à travers dix benchmarks (incluant GQA, MME, VQA-v2, etc.) :

Rétention de Performance : Sur LLaVA-NeXT-7B, MAP conserve 97,5 % de la performance du modèle non élagué tout en utilisant seulement 5,56 % des jetons visuels (160 jetons conservés sur 2880).
Accélération : À ce ratio de rétention, MAP atteint une accélération de 7,44× lors de la phase de prefill du LLM et une accélération de bout en bout de 3,09×.
Comparaison : MAP surpasse systématiquement les méthodes de pointe (ex: LearnPruner, CDPruner, ZOO-Prune) à travers divers ratios de rétention de jetons. L'écart de performance s'élargit à mesure que le ratio de rétention diminue, indiquant la robustesse de MAP sous une compression agressive.
Études d'Ablation :
- Sélection de l'Enseignant : QCTS surpasse la sélection de couche fixe (ex: Couche 14) et la fusion multi-couches, confirmant la nécessité d'une sélection d'enseignant spécifique à l'échantillon.
- Diversité : La sélection sensible à la diversité apporte des gains marginaux pour les budgets de jetons élevés, mais améliore significativement les performances lors d'un élagage agressif (ex: 32 jetons), atténuant la redondance parmi les jetons à score élevé.
- Efficacité : MAP réduit l'utilisation du cache KV jusqu'à 91,39 % sur LLaVA-NeXT-7B.

Signification

L'article affirme que MAP répond au compromis fondamental entre l'efficacité de l'élagage et le coût de l'inférence. En prédisant le signal d'attention optimal plutôt qu'en le calculant, MAP permet un élagage précoce des jetons qui préserve les preuves visuelles pertinentes à la question sans la surcharge de calcul liée à la traversée de multiples couches du LLM. La méthode est présentée comme une solution généralisable qui améliore l'efficacité des MLLM tout en maintenant une haute précision, particulièrement dans les scénarios nécessitant une réduction agressive de jetons.