Can Cross-Layer Transcoders Replace Vision Transformer Activations? An Interpretable Perspective on Vision

1. Problématique et Contexte

Les modèles de fondation, en particulier les Vision Transformers (ViT), ont démontré des performances exceptionnelles dans diverses tâches de vision par ordinateur. Cependant, leur fonctionnement interne reste largement opaque ("boîte noire"), ce qui limite leur interprétabilité, leur contrôlabilité et leur fiabilité.

Les approches existantes pour l'interprétabilité, telles que les Autoencodeurs Creux (Sparse Autoencoders - SAE), se concentrent sur l'extraction de caractéristiques interprétables au sein d'une seule couche à la fois. Cette limitation locale empêche de capturer la structure de calcul inter-couches (cross-layer) et la contribution relative de chaque couche dans la formation de la représentation finale. Il manque donc une méthode capable de décomposer fidèlement comment l'information est transformée à travers la profondeur du réseau et d'attribuer la responsabilité de la décision finale à des couches spécifiques.

2. Méthodologie : Les Transcodeurs Inter-couches (CLT)

Les auteurs proposent l'adoption des Transcodeurs Inter-couches (Cross-Layer Transcoders - CLT) comme modèles proxy fiables, creux et conscients de la profondeur pour les blocs MLP (Multi-Layer Perceptron) des ViT.

Architecture et Fonctionnement

Le CLT utilise un schéma encodeur-décodeur pour reconstruire les activations post-MLP d'une couche cible à partir des caractéristiques creuses (sparse embeddings) apprises des couches précédentes.

• Encodage : Pour chaque couche $i$, un encodeur linéaire apprend à mapper les activations pré-MLP ($x_i$) vers un vecteur de caractéristiques creuses ($z_i$). Trois fonctions d'activation non linéaires sont explorées pour générer ces codes : JumpReLU, ReLU-Top-k et Abs-Top-k.
• Transport Inter-couche : Pour reconstruire la sortie d'une couche cible $j$, le CLT agrège les contributions de toutes les couches sources $i \le j$. La reconstruction $\hat{y}_j$ est une somme additive :
  $$\hat{y}_j = \sum_{i=1}^{j} z_i W_{dec}^{i \to j}$$
  où $W_{dec}^{i \to j}$ est une matrice de décodage apprise spécifiquement pour la contribution de la couche $i$ vers la couche $j$.
• Structure Triangulaire : Cette architecture impose une structure causale stricte : une couche ne peut être reconstruite qu'à partir des couches qui la précèdent, évitant ainsi toute fuite d'information depuis les couches futures.

Entraînement et Remplacement

Les CLT sont entraînés sur des modèles ViT gelés (CLIP ViT-B/32 et ViT-B/16) pour minimiser l'erreur de reconstruction entre les sorties MLP originales et les sorties reconstruites, tout en pénalisant la densité des codes.
Une fois entraînés, les CLT peuvent remplacer fonctionnellement les blocs MLP originaux lors de l'inférence (modèle de remplacement), permettant d'analyser le modèle original sans compromettre ses performances sur les tâches en aval.

3. Contributions Clés

1. Remplacement Fonctionnel des MLP : Les auteurs démontrent que les CLT peuvent remplacer les blocs MLP, en particulier dans les couches tardives ou pour le token global [CLS], tout en préservant, voire en améliorant dans certains cas, la précision de classification zero-shot.
2. Attribution Inter-couche Fiable : La structure linéaire du CLT permet de calculer des scores de contribution qui quantifient l'influence de chaque couche source sur la représentation d'une couche cible. Cela offre une interprétabilité de processus (process-level interpretability) inédite.
3. Découverte d'Asymétrie Structurelle : L'analyse révèle une différence fondamentale entre les types de tokens :
   • Les tokens de patch (spatiaux) montrent une attribution diagonale dominante : chaque couche explique principalement sa propre sortie, avec une influence limitée des autres couches.
   • Le token [CLS] (global) montre une attribution largement distribuée : la représentation finale est une agrégation de signaux sémantiques provenant de nombreuses couches précédentes, et non seulement de la dernière.
4. Validation par Ablation : Des expériences montrent que la représentation finale est concentrée sur un petit sous-ensemble de couches dominantes. Retenir seulement les 4 couches les plus importantes permet de récupérer la précision originale, tandis que supprimer la couche la plus influente dégrade fortement les performances.

4. Résultats Expérimentaux

Les expériences ont été menées sur trois jeux de données (CIFAR-100, COCO, ImageNet-100) avec deux architectures ViT (B/32 et B/16).

• Fidélité de Reconstruction : Les CLT atteignent une haute fidélité de reconstruction (similarité cosinus entre 0,92 et 0,97, $R^2$ entre 0,89 et 0,95) pour les activations post-MLP. ViT-B/16 est plus facile à approximer que ViT-B/32, probablement en raison d'une granularité de patch plus fine.
• Performance de Classification (Remplacement) :
  • Le remplacement des couches tardives (ex: 10→11) ou uniquement du token [CLS] préserve la précision zero-shot (parfois même une légère amélioration).
  • Le remplacement complet de toutes les couches et de tous les tokens entraîne une dégradation, surtout pour les tokens de patch, soulignant la difficulté de remplacer les calculs des premières couches.
• Interprétabilité :
  • Les cartes de chaleur de contribution confirment que le token [CLS] intègre l'information de manière transversale, tandis que les patchs restent locaux.
  • Les ablations basées sur les scores de projection confirment que ces scores sont "fidèles" : supprimer les couches à fort score dégrade le modèle, et garder les meilleurs scores suffit à maintenir la performance.
• Visualisation : Un cadre de récupération basé sur les codes creux permet de visualiser comment les représentations évoluent sémantiquement à travers les couches (des statistiques de couleur/lumière aux concepts sémantiques complexes).

5. Signification et Conclusion

Cet article établit les Transcodeurs Inter-couches (CLT) comme un cadre principiel pour l'interprétabilité structurée des modèles de vision.

• Avantage par rapport aux SAE : Contrairement aux SAE qui opèrent localement, les CLT capturent la dynamique inter-couche, révélant comment les couches antérieures contribuent causalement à la représentation finale.
• Proxy Fiable : La capacité à remplacer les MLP sans perte de performance valide les CLT comme des substituts fiables pour l'analyse interne des modèles.
• Insight Architectural : La découverte de l'asymétrie entre la nature locale des patchs et la nature globale/intégrée du token [CLS] offre une nouvelle compréhension de la façon dont les ViT traitent l'information.

En conclusion, les CLT offrent une voie prometteuse pour rendre les modèles de vision plus transparents, contrôlables et dignes de confiance, en passant d'une analyse statique des couches à une analyse dynamique et causale de leur interaction.