Résumé technique : Cacher pour voir : Masquage des préfixes de raisonnement pour la pensée ancrée visuellement dans la distillation de modèles vision-langage

Énoncé du problème

Les approches récentes « réfléchir-répondre » dans les modèles vision-langage (VLM), telles que Qwen3-VL-Thinking, améliorent considérablement les performances de raisonnement en générant des étapes de raisonnement intermédiaires avant la réponse finale. Cependant, ces longues traces de raisonnement entraînent des coûts computationnels substantiels, limitant leur déploiement dans des scénarios aux ressources contraintes. La distillation de connaissances (KD) offre une solution en transférant les capacités de grands modèles enseignants vers des modèles étudiants compacts.

Un défi critique dans la distillation des VLM « réfléchir-répondre » est l'oubli visuel. Au fur et à mesure que le modèle étudiant génère une longue trace de raisonnement, il tend à s'appuyer fortement sur les préfixes de raisonnement textuels accumulés (les tokens précédemment générés) pour prédire les tokens suivants. Ces préfixes contiennent souvent suffisamment d'informations pour imiter la trajectoire de l'enseignant sans que l'étudiant ait besoin de réexaminer l'entrée visuelle. Par conséquent, l'étudiant apprend à ignorer les preuves visuelles, entraînant une dégradation des capacités de raisonnement ancré visuellement. Les méthodes de distillation existantes qui transfèrent des cartes d'attention visuelle ou des relations entre tokens échouent souvent à adresser cette tendance spécifique à prendre des raccourcis via des indices textuels.

Méthodologie : Masking-KD

Les auteurs proposent Masking-KD, un cadre de distillation novateur conçu pour forcer le modèle étudiant à ancrer son raisonnement dans des preuves visuelles en obscurcissant délibérément les indices textuels les plus influents.

Mécanisme central

Au lieu d'utiliser le masque causal standard (qui permet à l'étudiant de voir tous les tokens précédents), Masking-KD emploie un masque de préfixe de raisonnement saillant. Ce masque bloque l'accès aux tokens futurs (causalité standard) et à des préfixes de raisonnement saillants spécifiques sur lesquels l'étudiant s'appuierait autrement pour imiter l'enseignant. En supprimant ces indices textuels à haute influence, l'étudiant est contraint d'utiliser des preuves visuelles pour compenser les informations manquantes.

Le cadre se compose de deux composants adaptatifs :

1. Masquage des préfixes de raisonnement saillants au niveau du token :

   • Identification : Avant l'étape principale de distillation, un passage avant auxiliaire du modèle étudiant (sous un masque causal standard) est effectué. Cela génère une carte d'attention réponse-à-réponse ($A^{resp}$), qui révèle quels tokens précédents l'étudiant consulte le plus intensément lors de la prédiction du token actuel.
   • Sélection : Pour chaque prédiction de token suivant, la méthode identifie les $\rho$ préfixes les plus saillants en fonction de leurs poids d'attention. Ce sont les tokens fournissant le « raccourci » le plus fort vers la réponse.
   • Masquage : Ces préfixes saillants identifiés sont masqués (réglés à $-\infty$ dans les logits d'attention) lors de l'étape réelle d'entraînement par distillation, empêchant l'étudiant de s'y attarder.

2. Planification adaptative du budget de masquage auto-rythmée :

   • La difficulté de distiller chaque token varie. Les tokens facilement imités (faible divergence entre l'enseignant et l'étudiant) fournissent des signaux d'apprentissage faibles si l'étudiant peut simplement copier le préfixe.
   • La méthode calcule la divergence KL inverse au niveau du token entre les distributions de l'enseignant et de l'étudiant lors du passage avant auxiliaire.
   • Un seuil de ratio cumulatif ($\rho_n$) est attribué de manière adaptative à chaque token en fonction de cette difficulté. Les tokens facilement imités (faible divergence KL) reçoivent un budget de masquage plus élevé (davantage de préfixes sont masqués) pour forcer l'étudiant à s'appuyer sur des indices visuels. Les tokens plus difficiles à apprendre reçoivent moins de masquage pour assurer un entraînement stable.

Fonction objectif

L'étudiant est entraîné pour minimiser la divergence KL inverse entre sa distribution prédictive (dans le contexte masqué) et la distribution de l'enseignant (dans le contexte complet). La fonction de perte est définie comme suit :
$$\mathcal{L}_{Distill} = \frac{1}{N} \sum_{n=1}^{N} \sum_{y \in \mathcal{V}} p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}}) \log \frac{p_s(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \tilde{\mathbf{M}})}{p_t(y \mid \mathbf{x}_v, \mathbf{x}_q, \mathbf{y}_{<n}, \mathbf{M})}$$
Où $\tilde{\mathbf{M}}$ est le masque de préfixe de raisonnement saillant et $\mathbf{M}$ est le masque causal standard.

Contributions clés

1. Cadre Masking-KD : Un nouveau paradigme de distillation qui améliore la pensée ancrée visuellement en masquant les préfixes de raisonnement saillants, forçant l'étudiant à utiliser des preuves visuelles plutôt que de s'appuyer sur des raccourcis textuels.
2. Masquage adaptatif au niveau du token : Un mécanisme qui identifie et masque dynamiquement les préfixes de raisonnement les plus influents pour chaque prédiction de token spécifique, reconnaissant que le contexte critique varie au cours des étapes de décodage.
3. Planification de budget auto-rythmée : Une stratégie adaptative qui ajuste l'intensité du masquage en fonction de la difficulté de distillation (divergence KL) de chaque token, amplifiant les signaux d'apprentissage pour les tokens facilement imités.
4. Validation empirique : Des expériences extensives démontrant que cette approche surpasse les méthodes de distillation VLM existantes et les bases de référence de l'auto-distillation.

Résultats expérimentaux

Les auteurs ont évalué Masking-KD sur des benchmarks de raisonnement multimodal, notamment Geometry-3K, MathVista, We-Math, MMK12, MathVerse, LogicVista et MMMU-Pro.

• Gains de performance : Masking-KD a obtenu des résultats de l'état de l'art sur toutes les tailles de modèles (2B, 4B et 8B). Notamment, le modèle 4B distillé a surpassé le modèle enseignant 8B non distillé, et le modèle 2B a surpassé le modèle enseignant 4B non distillé.
• Comparaison avec les bases de référence : La méthode a nettement surpassé la distillation standard au niveau de la réponse, ainsi que les méthodes transférant explicitement des motifs visuels (par exemple, LLaVA-KD, CompoDistill, Align-TI).
• Auto-distillation : Le cadre s'est avéré efficace dans des configurations d'auto-distillation (où l'étudiant sert de son propre enseignant), surpassant d'autres méthodes d'auto-distillation comme OPSD.
• Analyse de l'attention visuelle : Des analyses qualitatives et quantitatives ont confirmé que les modèles Masking-KD maintiennent des ratios d'attention visuelle plus élevés tout au long du processus de génération par rapport aux bases de référence, atténuant efficacement le phénomène d'« oubli visuel ». Il a été observé que les modèles activaient plus fortement les régions d'image pertinentes lorsque les indices textuels étaient masqués.

Importance et affirmations

L'article affirme que le principal goulot d'étranglement dans la distillation des VLM « réfléchir-répondre » est la tendance de l'étudiant à exploiter les raccourcis textuels fournis par les préfixes de raisonnement accumulés, négligeant ainsi les entrées visuelles.

L'importance de Masking-KD réside dans son passage du transfert explicite des connaissances visuelles de l'enseignant (par exemple, via des cartes d'attention) à la guidance de l'étudiant pour développer son propre processus de raisonnement ancré visuellement. En « cachant » les indices textuels les plus évidents, la méthode force l'étudiant à « voir » l'image pour résoudre le problème. Les auteurs concluent que cette approche est plus efficace pour la distillation du raisonnement VLM que les méthodes précédentes se concentrant uniquement sur l'imitation de motifs visuels internes, car elle adresse directement la cause racine de l'oubli visuel dans les tâches de raisonnement à longue chaîne.