From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

Résumé Technique : De la Sparsité à la Simplicité (S2S)

Énoncé du Problème
Les mécanismes d'attention auto-attentionnelle constituent le fondement du préentraînement des transformers à grande échelle, permettant un mélange flexible des tokens et une modélisation des dépendances à long terme. Cependant, leur complexité quadratique par rapport à la longueur de la séquence entraîne une latence d'inférence élevée et une surcharge mémoire. Bien que le remplacement de l'attention par des modules séquentiels plus simples (par exemple, les RNN, les modèles d'espace d'état) soit une voie attrayante pour l'efficacité, une substitution naïve conduit souvent à une dégradation significative de la précision, en particulier à plus grande échelle. Les approches d'efficacité existantes modifient souvent le calcul de l'attention lui-même ou élaguent les tokens, mais conservent l'attention comme opérateur central. Cet article aborde le défi du remplacement complet des couches d'attention par des modules séquentiels sans perte substantielle de performance.

Méthodologie
L'article propose un cadre appelé Sparsity to Simplicity (S2S), qui postule que les transformers préentraînés décomposent les dépendances complexes des tokens en des mappings par blocs de complexité effective variable. Les auteurs émettent l'hypothèse que les couches présentant des motifs d'attention clairsemés sont intrinsèquement plus faciles à approximer avec des modules séquentiels plus simples.

La méthodologie se compose de trois composantes principales :

1. Remplacement Unifié des Couches d'Attention :

   • Les auteurs introduisent un cadre « plug-and-play » où des couches d'attention spécifiques dans un Vision Transformer (ViT) préentraîné sont remplacées par des mélangeurs de tokens séquentiels (spécifiquement BiMamba et BiLSTM), tout en préservant l'interface de la couche originale (LayerNorm, connexions résiduelles et MLP).
   • Les modules séquentiels adoptent une structure bidirectionnelle multi-têtes pour correspondre à la nature non causale et à contexte complet de l'attention ViT.
   • Une supervision alignée par couche est utilisée lors de la distillation : la sortie du mélangeur séquentiel remplacé est comparée directement à la sortie de la couche d'attention originale à la même profondeur en utilisant une perte d'erreur quadratique moyenne (MSE).

2. Distillation Guidée par la Sparsité :

   • Pour contrôler et étudier l'impact de la sparsité, les auteurs intègrent A-ViT (Adaptive ViT), qui utilise des scores d'arrêt pour générer des masques binaires de rétention de tokens.
   • Étape 1 (Distillation) : Le modèle étudiant (avec remplacements séquentiels) est entraîné sous les mêmes masques de tokens par couche générés par l'enseignant (A-ViT). Cela garantit que l'étudiant et l'enseignant opèrent sur les mêmes tokens actifs, isolant la difficulté d'approximer la fonction de mélange des tokens des différences de routage des tokens. La perte inclut une perte de similarité, un terme d'alignement d'arrêt et une perte de classification.
   • Étape 2 (Affinage) : L'étudiant est affiné en utilisant ses propres modules d'arrêt appris pour générer des masques, lui permettant d'opérer de manière autonome tout en maintenant la performance de la tâche.

3. Stratégie de Remplacement Sparse Sensible à la Couche :

   • Au lieu d'un remplacement uniforme, la stratégie cible des groupes de couches spécifiques basés sur leurs profils de sparsité naturels. Les auteurs ne remplacent que les couches les plus profondes (qui présentent la sparsité naturelle la plus élevée) et appliquent le cadre de distillation guidée par la sparsité à ces blocs spécifiques.

Contributions Clés

• L'Hypothèse S2S : L'article établit que la sparsité de l'attention est une signature de la complexité du mapping. Les blocs clairsemés dans les transformers préentraînés peuvent être remplacés par des modules séquentiels plus simples avec une perte minimale, tandis que les blocs denses sont plus difficiles à approximer.
• Cadre de Remplacement Contrôlé : Une méthode unifiée pour échanger l'attention contre des modules séquentiels (Mamba/LSTM) tout en maintenant la compatibilité des interfaces et en permettant une analyse sélective par couche.
• Contrôle Explicite de la Sparsité : L'introduction d'un pipeline de distillation où la sparsité induite par l'enseignant (via les masques A-ViT) est explicitement alignée avec l'étudiant, démontrant que l'augmentation de la sparsité de l'enseignant réduit systématiquement l'écart de performance entre l'enseignant et l'étudiant.
• Compromis Optimisé : Une stratégie qui combine la sparsité naturelle par couche avec une sparsification explicite pour atteindre de meilleurs compromis précision-efficacité que les bases d'attention clairsemée.

Résultats Expérimentaux
Des expériences ont été menées sur ImageNet en utilisant des architectures DeiT (Tiny, Small, Base) avec BiMamba et BiLSTM comme substituts.

• Sparsité Naturelle et Remplaçabilité : Les visualisations des cartes d'interaction des tokens montrent que, tandis que les couches peu profondes présentent un mélange global dense, les couches plus profondes deviennent nettement plus clairsemées et plus sélectives. Des expériences contrôlées remplaçant différents groupes de couches révèlent que le remplacement des couches plus profondes (plus clairsemées) entraîne une baisse de précision significativement plus faible que le remplacement des couches plus superficielles (plus denses).
• Effet de la Sparsité Explicite : Lorsque le ratio de rétention des tokens est réduit via A-ViT, l'écart de précision entre l'enseignant basé sur l'attention et l'étudiant séquentiel rétrécit systématiquement. Cela confirme que la sparsification explicite simplifie le mapping cible, facilitant l'apprentissage pour les modules séquentiels.
• Compromis Précision-Efficacité : Le remplacement sparse sensible à la couche proposé (remplacement des 3 couches les plus profondes uniquement par BiMamba sous des budgets de sparsité appariés) atteint :
  • Une précision comparable ou légèrement supérieure aux bases d'attention clairsemée (A-ViT) à des taux de rétention de tokens similaires.
  • Des gains d'efficacité significatifs : Jusqu'à 1,71× d'accélération estimée à une résolution de 384² par rapport à la base dense, avec un débit de mélange de tokens supérieur à celui des modèles d'attention clairsemée pour des nombres de tokens plus élevés.

Signification et Revendications
L'article revendique que le remplacement de l'attention doit être traité comme un problème structuré et dépendant de la couche plutôt que comme un échange d'architecture uniforme. La signification de ce travail réside dans la démonstration que :

1. La sparsité est un signal de remplaçabilité : Les modèles préentraînés décomposent naturellement la complexité de telle sorte que les couches clairsemées sont les plus propices à la simplification.
2. La sparsification explicite aide la distillation : Imposer la sparsité à l'enseignant rend la tâche d'apprentissage pour l'étudiant séquentiel plus facile, réduisant l'écart de performance.
3. Le remplacement stratégique surpasse les approches uniformes : En concentrant le remplacement sur des couches naturellement clairsemées et en guidant le processus avec une sparsité contrôlée, on peut atteindre un équilibre supérieur entre précision et latence par rapport au simple élagage de l'attention ou au remplacement complet du modèle.

Les auteurs concluent que cette perspective encourage les recherches futures sur des modèles hybrides qui exploitent les connaissances des transformers préentraînés tout en intégrant des modules de séquence plus simples et plus efficaces là où la structure interne du modèle le permet.