Récemment, les recherches de l'équipe Mamba ont été remarquables : des chercheurs d'universités comme Cornell et Princeton ont réussi à « distiller » le grand modèle Transformer Llama en Mamba, et ont conçu un nouvel algorithme de décodage d'inférence, améliorant considérablement la vitesse d'inférence du modèle.
L'objectif des chercheurs était de transformer Llama en Mamba. Pourquoi ? Car entraîner un grand modèle à partir de zéro est coûteux, et bien que Mamba ait suscité un intérêt considérable depuis son apparition, peu d'équipes entraînent elles-mêmes des modèles Mamba à grande échelle. Bien que des variantes réputées existent, comme Jamba d'AI21 et Hybrid Mamba2 de NVIDIA, de nombreuses connaissances sont contenues dans les modèles Transformer performants. Si nous pouvions capturer ces connaissances et affiner le Transformer en Mamba, le problème serait résolu.
L'équipe de recherche a combiné plusieurs méthodes, notamment la distillation progressive, le réglage fin supervisé et l'optimisation des préférences dirigées, pour atteindre cet objectif. Il est important de noter que la vitesse est essentielle, sans compromettre les performances. Mamba présente un avantage significatif dans l'inférence de longues séquences, et les modèles Transformer disposent également de solutions d'accélération de l'inférence, comme le décodage spéculatif. Comme la structure unique de Mamba ne permet pas d'appliquer directement ces solutions, les chercheurs ont spécialement conçu un nouvel algorithme et l'ont combiné aux caractéristiques matérielles pour implémenter le décodage spéculatif basé sur Mamba.
Finalement, les chercheurs ont réussi à convertir Zephyr-7B et Llama-38B en modèles RNN linéaires, avec des performances comparables aux modèles standard avant distillation. L'ensemble du processus d'entraînement n'a utilisé que 20 milliards de jetons, des résultats comparables à ceux du modèle Mamba7B entraîné à partir de zéro avec 1,2 T de jetons et du modèle NVIDIA Hybrid Mamba2 entraîné avec 3,5 T de jetons.
Sur le plan technique, les RNN linéaires et l'attention linéaire sont similaires, permettant aux chercheurs de réutiliser directement les matrices de projection du mécanisme d'attention et de construire le modèle via l'initialisation des paramètres. De plus, l'équipe de recherche a gelé les paramètres des couches MLP du Transformer, remplaçant progressivement les têtes d'attention par des couches RNN linéaires (c'est-à-dire Mamba), et a traité l'attention de requête groupée partageant les clés et les valeurs entre les têtes.
Au cours du processus de distillation, une stratégie de remplacement progressif des couches d'attention a été adoptée. Le réglage fin supervisé comprend deux méthodes principales : une basée sur la divergence de Kullback-Leibler (KL) au niveau des mots, et une autre sur la distillation des connaissances au niveau des séquences. Pour l'étape d'optimisation des préférences de l'utilisateur, l'équipe a utilisé l'optimisation des préférences directes (DPO), en comparant avec la sortie du modèle enseignant pour garantir que le modèle génère un contenu mieux adapté aux attentes de l'utilisateur.
Ensuite, les chercheurs ont commencé à appliquer le décodage spéculatif des Transformers au modèle Mamba. Le décodage spéculatif peut être simplement compris comme l'utilisation d'un petit modèle pour générer plusieurs sorties, puis la vérification de ces sorties par un grand modèle. Le petit modèle fonctionne rapidement et peut générer rapidement plusieurs vecteurs de sortie, tandis que le grand modèle évalue la précision de ces sorties, améliorant ainsi la vitesse d'inférence globale.
Pour réaliser ce processus, les chercheurs ont conçu un algorithme qui utilise à chaque fois un petit modèle pour générer K sorties brutes, puis le grand modèle renvoie la sortie finale et le cache des états intermédiaires après vérification. Cette méthode a donné de bons résultats sur les GPU, Mamba2.8B atteignant une accélération de l'inférence de 1,5 fois et un taux d'acceptation de 60 %. Bien que les résultats varient selon les architectures de GPU, l'équipe de recherche a optimisé davantage les résultats en fusionnant les noyaux et en ajustant les méthodes d'implémentation.
Lors de la phase expérimentale, les chercheurs ont utilisé Zephyr-7B et Llama-3Instruct8B pour un entraînement de distillation en trois étapes, réussissant finalement à reproduire les résultats de la recherche en seulement 3 à 4 jours sur 8 cartes A100 de 80 Go.
Cette recherche montre non seulement le chemin de la transformation de Llama en Mamba, mais offre également de nouvelles pistes pour améliorer la vitesse et les performances des modèles futurs.
Adresse de l'article : https://arxiv.org/pdf/2408.15237
