Steering Sparse Autoencoder Latents to Control Dynamic Head Pruning in Vision Transformers (Student Abstract)

1. Problématique

Les Vision Transformers (ViT) utilisent l'attention multi-têtes pour capturer des interactions complexes entre les tokens. Cependant, une grande partie de ces têtes d'attention est redondante, entraînant une augmentation du coût computationnel sans gain proportionnel de performance.
Bien que des cadres adaptatifs comme AdaViT aient été développés pour élaguer dynamiquement les têtes en fonction de l'entrée (en utilisant des réseaux de décision basés sur les embeddings résiduels), ces méthodes souffrent d'un manque crucial :

• Opacité : Les politiques d'élagage sont difficiles à interpréter car elles reposent sur des embeddings denses et poly-sémantiques.
• Manque de contrôle : Il est impossible de manipuler ou d'expliquer les décisions d'élagage au niveau latent.
  L'objectif de ce travail est de combler ce fossé en rendant l'élagage dynamique à la fois interprétable et contrôlable au niveau des caractéristiques latentes.

2. Méthodologie

Les auteurs proposent un cadre novateur intégrant des Autoencodeurs Parcimonieux (Sparse Autoencoders - SAE) avec l'élagage dynamique de têtes dans les ViT. La méthodologie se décompose en trois étapes principales :

A. Entraînement de la base (AdaViT)

• Un ViT (avec des réseaux de décision par couche) est entraîné conjointement pour sélectionner les têtes à conserver.
• Le réseau de décision prend le token CLS (représentant le contexte global) du flux résiduel final et produit des logits d'importance pour chaque tête.
• Un masque binaire est généré via un échantillonnage Gumbel-Sigmoid pour éliminer les têtes redondantes tout en maintenant une précision cible.

B. Entraînement du SAE

• Une fois le modèle ViT entraîné, le token CLS de la dernière couche ($x \in \mathbb{R}^d$) est extrait.
• Un k-sparse autoencoder est entraîné sur ces embeddings pour les décomposer en caractéristiques latentes parcimonieuses ($z$).
• Architecture : L'encodeur projette l'entrée dans un espace latent de plus grande dimension ($n=3072$ pour $d=384$, soit un facteur 8x) avec une activation Top-K ($k=64$). Cela force le modèle à ne garder que les $k$ dimensions les plus actives, favorisant ainsi des concepts mono-sémantiques et interprétables.
• L'objectif est de minimiser l'erreur quadratique moyenne (MSE) entre l'entrée et la reconstruction.

C. Guidage (Steering) et Contrôle de l'Élagage

C'est l'étape clé de l'approche. Pour influencer les décisions d'élagage :

1. Amplification : Des dimensions latentes spécifiques ($S$) sont sélectionnées et amplifiées par une constante $\alpha$ : $z'_i = z_i + \alpha$ si $i \in S$.
2. Reconstruction : Le vecteur latent guidé $\tilde{z}$ est reconstruit en un embedding guidé $\tilde{x}$.
3. Décision : Cet embedding $\tilde{x}$ est réinjecté dans le réseau de décision du ViT pour générer un nouveau masque d'élagage.

Trois stratégies de sélection des dimensions $S$ sont testées :

• Par classe (Per-class) : Les latents les plus fréquemment activés pour une classe spécifique.
• Global : Les latents les plus fréquents sur toutes les classes.
• Aléatoire.

3. Contributions Clés

• Intégration SAE-Élagage : Première approche combinant les SAE pour le guidage (steering) et l'élagage dynamique des têtes dans les ViT.
• Interprétabilité Mécanistique : Démontre que les latents parcimonieux peuvent révéler des sous-ensembles de têtes spécifiques à une classe, rendant le processus de décision transparent.
• Contrôle Granulaire : Permet de modifier le comportement d'élagage (efficacité vs précision) en manipulant simplement des dimensions latentes spécifiques, sans réentraîner le modèle ViT.

4. Résultats Expérimentaux

Les expériences ont été menées sur un ViT-Small pré-entraîné sur ImageNet, finetuné sur CIFAR-100.

• Performance du SAE : Le SAE préserve l'information essentielle (la reconstruction de l'embedding n'impacte que très légèrement la précision : -0,12 %).
• Efficacité du Guidage :
  • Le guidage par classe (Per-class steering) s'avère supérieur. Il réduit significativement l'utilisation des têtes tout en préservant, voire en améliorant, la précision.
  • Les stratégies globales ou aléatoires entraînent une chute de précision plus marquée lorsque l'amplification ($\alpha$) augmente.
• Cas Concrets (Exemples) :
  • Pour la classe "bowl" (bol) : La précision passe de 76 % à 82 % tandis que l'utilisation des têtes chute de 0,72 à 0,33. Le modèle se concentre principalement sur les têtes h2 et h5.
  • Pour la classe "pine tree" (pin) : Précision de 79 % à 84 %, utilisation des têtes de 0,93 à 0,35 (têtes h2 et h3).
• Découvertes Sémantiques : Des classes sémantiquement liées (ex: "bowl" et "plate") partagent des sous-ensembles de têtes similaires (h2, h5), révélant que le guidage par classe expose des relations sémantiques au niveau des têtes d'attention.

5. Signification et Conclusion

Ce travail démontre que les caractéristiques latentes parcimonieuses extraites par un SAE offrent un moyen efficace d'interpréter et de contrôler l'élagage dynamique dans les ViT.

• Bridging the Gap : L'étude comble le fossé entre l'efficacité computationnelle (réduction des têtes) et l'interprétabilité mécanique.
• Contrôle Prédictif : Elle permet de concevoir des politiques d'élagage spécifiques à une tâche ou une classe, optimisant le compromis coût/précision.
• Perspectives : Bien que l'étude actuelle se limite à la dernière couche et à de petits jeux de données, les auteurs prévoient d'étendre ce cadre aux couches précédentes et aux modèles de fondation (foundation models) pour une applicabilité plus large.

En résumé, cette approche transforme l'élagage dynamique d'une boîte noire opaque en un processus transparent et pilotable, ouvrant la voie à des modèles de vision plus efficaces et compréhensibles.