💬 NLP

Attention to Mamba: A Recipe for Cross-Architecture Distillation

Cet article propose une méthode de distillation en deux étapes, combinant l'adaptation de l'astuce du noyau et une initialisation rigoureuse, pour transférer efficacement les connaissances d'un Transformer vers une architecture Mamba pure sans blocs d'attention, préservant ainsi les performances du modèle enseignant.

Auteurs originaux : Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodríguez, Luca Zappella, Federico Danieli

Publié 2026-04-17

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodríguez, Luca Zappella, Federico Danieli

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 avez un chef étoilé (le modèle Transformer) qui cuisine des plats délicieux (le langage humain). Ce chef est incroyable, mais il a un gros défaut : il utilise une technique de cuisson très lente et gourmande en énergie (l'attention quadratique). Pour chaque nouveau plat, il doit regarder tous les ingrédients précédents un par un, ce qui prend beaucoup de temps et d'électricité.

D'un autre côté, il existe un nouveau type de cuisinier, le Mamba, qui est ultra-rapide et économe en énergie. Il peut cuisiner en regardant les ingrédients de manière linéaire, comme une chaîne de montage. Mais il y a un problème : le Mamba est un apprenti qui n'a pas encore la finesse du chef étoilé. Ses plats sont rapides, mais souvent moins bons.

L'objectif de ce papier est de répondre à une question simple : Comment apprendre au Mamba à cuisiner aussi bien que le Chef Étoilé, sans avoir à le former depuis zéro pendant des années ?

Voici la recette magique découverte par les chercheurs, expliquée simplement :

1. Le Problème : On ne peut pas copier-coller directement

Si vous essayez de donner directement les recettes du Chef Étoilé au Mamba (ce qu'on appelle une "distillation naïve"), ça échoue. C'est comme essayer d'enseigner la cuisine française à un robot conçu pour faire de la pizza : ils ne parlent pas le même langage. Le robot se perd et produit des résultats médiocres.

2. La Solution : Le "Pont" en deux étapes

Au lieu de forcer le Mamba à comprendre le Chef directement, les chercheurs ont construit un pont en deux étapes. C'est comme si on traduisait d'abord le langage du Chef en un langage intermédiaire, avant de l'enseigner au Mamba.

Étape 1 : La Traduction (De "Complexe" à "Simplifié")

Le Chef Étoilé utilise une technique complexe appelée "Softmax" (une façon très précise mais lente de pondérer les ingrédients).

L'astuce : Les chercheurs ont d'abord appris au Mamba à utiliser une version simplifiée de cette technique, appelée Attention Linéaire (ou "Hedgehog").
L'analogie : Imaginez que le Chef dit : "Regarde tous les ingrédients pour décider du sel". La version simplifiée dit : "Regarde juste les ingrédients les plus importants d'une manière plus rapide".
À ce stade, le Mamba commence à comprendre la logique du Chef, mais avec une technique plus simple. C'est comme si on lui donnait un manuel de cuisine simplifié.

Étape 2 : L'Enrichissement (De "Simplifié" à "Expert")

Une fois que le Mamba a compris les bases via cette version simplifiée, on lui ajoute les outils spéciaux qui font sa force (la mémoire sélective, les convolutions, etc.).

L'astuce : On prend les connaissances acquises à l'étape 1 et on les utilise comme fondation pour entraîner le Mamba complet.
L'analogie : C'est comme si l'apprenti avait déjà lu le manuel simplifié. Maintenant, on lui donne les outils de pro (le robot de cuisine Mamba) et on lui dit : "Maintenant, utilise ces outils pour appliquer ce que tu as appris."
Résultat : Le Mamba garde la vitesse de l'apprenti, mais acquiert la finesse du Chef.

3. Les Résultats : Le Meilleur des deux mondes

Grâce à cette méthode en deux temps, les chercheurs ont créé un modèle (qu'ils appellent HedgeMamba) qui :

Est aussi rapide que le Mamba (il ne ralentit pas quand le texte est long).
Est aussi bon que le Chef Étoilé (il fait presque aussi bien sur les tests de compréhension et de raisonnement).

En chiffres, le modèle original (le Chef) avait un score de perplexité (une mesure d'erreur) de 13,86. Le nouveau modèle Mamba distillé a un score de 14,11. C'est une différence infime ! Ils ont réussi à transférer presque tout le savoir du Chef vers le robot rapide, en utilisant seulement une fraction de l'énergie habituelle.

En résumé

Ce papier nous dit : Ne jetez pas vos vieux modèles puissants (Transformers) juste parce qu'ils sont lents. Au lieu de les remplacer brutalement, utilisez cette "recette en deux étapes" pour transférer leur intelligence dans des modèles plus rapides (Mamba). C'est comme transformer une vieille voiture de course lente mais puissante en une voiture de course moderne, rapide et économe, sans perdre la puissance du moteur.

C'est une victoire pour l'efficacité : on garde la qualité, on gagne du temps et de l'argent.

1. Problématique et Motivation

Le domaine du traitement du langage naturel (NLP) a été dominé par les modèles basés sur le mécanisme d'Attention (Transformers), qui offrent d'excellentes performances mais souffrent d'une complexité computationnelle quadratique ( $O(L^2)$ ) par rapport à la longueur de la séquence $L$ . Cela entraîne des coûts de mémoire et d'inférence élevés.

Les Modèles à Espace d'État (SSM), comme Mamba, émergent comme une alternative prometteuse grâce à leur complexité linéaire ( $O(L)$ ) et leur efficacité à l'inférence. Cependant, les modèles SSM entraînés à partir de zéro peinent souvent à égaler les performances des Transformers pré-entraînés, surtout à grande échelle. De plus, la distillation directe (transfert de connaissances) d'un Transformer vers un Mamba s'est avérée inefficace, conduisant à une perte significative de performances.

L'objectif de ce travail est de définir une "recette" efficace pour distiller un modèle Transformer (enseignant) vers une architecture de type Mamba (élève), en préservant les performances du modèle enseignant tout en bénéficiant de l'efficacité computationnelle du SSM.

2. Méthodologie : Une Approche en Deux Étapes

Les auteurs proposent une méthode de distillation croisée structurée en deux étapes, fondée sur l'alignement architectural et une initialisation raisonnée. L'idée centrale est d'éviter le saut direct entre deux paradigmes computationnels radicalement différents en passant par une étape intermédiaire de linéarisation.

Étape 1 : De l'Attention Softmax à l'Attention Linéaire (Hedgehog)

La première étape consiste à remplacer le mécanisme d'Attention Softmax quadratique par une version linéaire, en utilisant la procédure Hedgehog (Zhang et al., 2024).

Principe : Au lieu de supprimer simplement l'opérateur Softmax (ce qui dégrade les performances), les auteurs apprennent une carte de caractéristiques (feature map) $\phi$ via un petit MLP.
Théorie : En s'appuyant sur le théorème de Mercer, l'opérateur exponentiel du Softmax est approximé par un produit scalaire de vecteurs de caractéristiques : $e^{x^Tx'} \approx \phi(x)^T\phi(x')$ .
Objectif : Transformer le module d'Attention du Transformer en une Attention Linéaire qui imite fidèlement le comportement du Softmax, tout en réduisant la complexité. Les poids de l'enseignant sont figés, sauf ceux de la carte de caractéristiques apprise.

Étape 2 : De l'Attention Linéaire à Mamba (HedgeMamba)

Une fois l'Attention linéarisée, elle sert d'initialisation pour le bloc Mamba.

Alignement des paramètres : Les auteurs établissent un lien mathématique direct entre l'Attention Linéaire et les SSM. Les matrices de l'Attention Linéaire ( $\hat{Q}, \hat{K}, \hat{V}$ $\hat{Q}, \hat{K}, \hat{V}$ ) sont mappées aux paramètres du SSM ( $C, B, X$ $C, B, X$ ).
- $B(X) \leftarrow \hat{K}(X) = \phi(K(X))$
- $C(X) \leftarrow \hat{Q}(X) = \phi(Q(X))$
- $X \leftarrow \hat{V}(X) = V(X)$
Initialisation : Les paramètres supplémentaires de Mamba (masque causal $\Lambda$ , convolution, branche de porte/gate) sont initialisés pour se comporter comme l'opérateur identité, assurant que le modèle Mamba démarre avec les performances de l'Attention Linéaire.
Normalisation : Une étape cruciale consiste à normaliser les scores d'attention dans le SSM pour imiter le comportement du Softmax, en dupliquant l'état caché et les entrées.
Affinage (Fine-tuning) : Le modèle complet (HedgeMamba) est ensuite affiné sur le même corpus (OpenWebText) en utilisant la perte d'entropie croisée par rapport aux vrais tokens, en débloquant progressivement les composants Mamba (convolution, porte).

3. Contributions Clés

Nouvelle recette de distillation croisée : Proposition d'une méthode en deux étapes (Attention $\to$ Attention Linéaire $\to$ Mamba) qui surmonte les échecs de la distillation directe.
Initialisation Principée : Démonstration que l'initialisation des paramètres Mamba via une carte de caractéristiques apprise (Hedgehog) est essentielle pour transférer efficacement les connaissances.
Architecture Hybride (HedgeMamba) : Conception d'un module de mélange de séquences qui combine la linéarisation Hedgehog avec les composants expressifs de Mamba (SSM, convolution, porte).
Analyse Empirique Rigoureuse : Études d'ablation, analyse de l'échelle (scaling) et analyse de sensibilité sur l'allocation du budget de tokens entre les deux étapes.

4. Résultats Expérimentaux

Les expériences ont été menées avec un modèle enseignant Pythia-1B et un budget de distillation de 10 milliards de tokens (soit environ 2,7 % du budget d'entraînement original du Pythia).

Performance Globale : Le modèle élève distillé (HedgeMamba) atteint une perplexité de 14,11, très proche de celle du modèle enseignant (13,86), surpassant largement la distillation directe (PPL > 100) et la baseline Hedgehog seule (PPL 14,89).
Tâches en aval : Sur une suite de benchmarks (ARC, BoolQ, Lambada, etc.), le modèle HedgeMamba conserve la majorité des performances du Teacher, obtenant des scores supérieurs à la baseline Hedgehog sur la plupart des tâches.
Analyse de l'allocation des tokens : L'étude de sensibilité révèle qu'une répartition 10 % pour l'étape 1 (apprentissage de la carte Hedgehog) et 90 % pour l'étape 2 (affinage Mamba) est optimale. Une répartition 50/50 (comme dans les travaux précédents) est sous-optimale.
Impact des composants : L'ajout de la branche de porte (gate branch) de Mamba s'est révélé être le composant ayant le plus grand impact positif sur la réduction de la perplexité et l'amélioration des tâches en aval.
Évolutivité : La méthode fonctionne également bien sur des modèles plus petits (160M, 410M), montrant une tendance cohérente d'amélioration par rapport aux approches linéaires pures.

5. Signification et Impact

Ce travail est significatif car il offre une voie pratique pour adopter les architectures SSM (comme Mamba) sans avoir à les entraîner à partir de zéro, ce qui est extrêmement coûteux.

Efficacité : Il permet de réduire la complexité d'inférence de quadratique à linéaire tout en conservant la qualité des modèles Transformer pré-entraînés.
Faisabilité : La méthode démontre qu'il est possible de distiller des modèles massifs avec un budget de tokens réduit, rendant l'adoption de SSMs accessible à des communautés avec moins de ressources computationnelles.
Fondement Théorique : En reliant mathématiquement l'Attention Linéaire et les SSM via une initialisation apprise, le papier comble le fossé théorique entre ces deux familles de modèles, suggérant que les SSMs peuvent hériter de la richesse sémantique des Transformers si l'initialisation est correctement conçue.

En résumé, les auteurs proposent une "recette" robuste pour transformer un Transformer en Mamba, prouvant que l'alignement architectural via une étape intermédiaire de linéarisation est la clé pour réussir la distillation entre ces architectures hétérogènes.