Retrofitting Linear Attention into Diffusion Language Models

Imaginez le monde de l'intelligence artificielle comme une immense et animée bibliothèque où un robot bibliothécaire super intelligent essaie d'écrire une histoire. Pendant des années, la méthode standard de travail de ce bibliothécaire était « autorégressive » : il écrivait un mot, s'arrêtait, réfléchissait à tout ce qu'il venait d'écrire, écrivait le mot suivant, s'arrêtait, et réfléchissait de nouveau. C'était comme construire une tour brique par brique, en attendant que la colle sèche avant d'ajouter la suivante. Bien que fiable, cette méthode était lente, surtout pour les histoires longues.

Récemment, un nouveau style appelé « Diffusion » est arrivé. Au lieu de construire brique par brique, le bibliothécaire de la Diffusion commence avec une page blanche remplie de points d'interrogation et essaie de deviner toute l'histoire d'un coup, en affinant ses suppositions encore et encore jusqu'à ce que les mots fassent sens. C'est comme avoir une équipe d'artistes peignant une fresque simultanément ; ils peuvent travailler beaucoup plus vite car ils n'attendent pas que la brique précédente sèche. Cependant, même cette méthode rapide présente un obstacle. À mesure que la fresque s'agrandit, les artistes doivent constamment regarder en arrière chaque mot qu'ils ont déjà décidé pour s'assurer que les nouveaux mots s'y accordent. C'est comme essayer de se souvenir d'un livre de 10 000 pages tout en écrivant la page 101, ce qui finit par ralentir tout le monde et remplit la mémoire de la bibliothèque.

C'est le puzzle que une équipe de chercheurs a abordé. Ils se sont demandé : « Pouvons-nous rendre cette méthode de Diffusion rapide encore plus rapide en changeant la façon dont le bibliothécaire se souvient du passé ? » Leur réponse est une astuce ingénieuse appelée « Attention Hybride par Blocs » (Block-Hybrid Attention). Ils ont réalisé que si le bibliothécaire a besoin de se souvenir parfaitement du paragraphe actuel (pour que les phrases coulent de source), il n'a pas besoin de relire chaque mot des chapitres précédents en haute définition. Au lieu de cela, il peut résumer les anciens chapitres en une petite « fiche de synthèse » de taille fixe. Cela permet au bibliothécaire d'écrire les nouvelles parties de l'histoire à la vitesse de l'éclair sans être freiné par le poids de l'histoire entière.

La Grande Idée : Un Système de Mémoire à Deux Vitesses

Le papier présente une méthode pour moderniser (ou mettre à niveau) un modèle de langage de diffusion existant et puissant (plus précisément un modèle de 16 milliards de paramètres appelé LLaDA 2.1) afin qu'il puisse utiliser cette stratégie de « fiche de synthèse » sans avoir besoin d'être réentraîné de zéro.

Considérez le cerveau du modèle comme ayant 20 couches différentes de réflexion. Dans le modèle original, chaque couche agit comme un bibliothécaire super attentif qui lit chaque mot précédent pour comprendre le mot actuel. C'est précis mais lourd. L'innovation des chercheurs, l'Attention Hybride par Blocs, divise le travail en deux modes :

1. Le Mode « Maintenant » (Attention Exacte) : Lorsque le modèle travaille sur le bloc de mots actuel qu'il est en train de générer, il conserve son mode « super attentif ». Il utilise la mémoire standard, très lourde, pour examiner chaque mot du paragraphe actuel afin de s'assurer que les phrases font parfaitement sens.
2. Le Mode « Passé » (Attention Linéaire) : Lorsque le modèle doit se souvenir de ce qui s'est passé dans les blocs précédents (les chapitres terminés), il passe à l'« Attention Linéaire ». Au lieu de relire tout le livre, il consulte un état de résumé de taille fixe. C'est comme consulter l'index d'un livre ou un résumé d'une page plutôt que de relire l'intégralité du texte. Ce résumé garde la même taille, peu importe la longueur du livre.

Comment ils ont fait : La Mise à Niveau en Deux Étapes

Les chercheurs n'ont pas seulement remplacé les pièces en espérant que tout fonctionne ; ils ont utilisé un processus rigoureux en deux étapes pour s'assurer que le modèle ne perde pas son intelligence lors de la mise à niveau.

• Étape 1 : L'Entraînement de l'Ombre. Ils ont pris le modèle original et intelligent (le « Enseignant ») et l'ont gelé. Ensuite, ils ont remplacé 6 des 20 couches d'attention par leur nouvelle version « Hybride par Blocs » (l' « Étudiant »). L'Étudiant a été entraîné pour imiter exactement la sortie de l'Enseignant, en utilisant une version « corrompue » du texte comme entrée. C'était comme avoir un étudiant qui observe l'ombre d'un chef cuisinier expert, essayant de copier ses mouvements exacts et de goûter les mêmes saveurs, mais uniquement pour les plats spécifiques que l'étudiant est chargé de modifier. Cette étape a duré environ 24 heures.
• Étape 2 : Le Fine-Tuning (Affinage). Une fois que l'Étudiant a appris à imiter l'Enseignant, les chercheurs ont laissé tout le modèle travailler à nouveau ensemble. Ils ont utilisé une technique appelée LoRA (Low-Rank Adaptation) pour effectuer des ajustements minuscules et précis sur les connexions du modèle. C'était comme donner un rapide « polissage » au modèle mis à jour pour corriger les petits défauts qui auraient pu apparaître parce que les nouvelles parties communiquaient désormais avec les anciennes. Cette étape a duré environ 6 heures.

Les Résultats : Plus Rapide, Plus Léger et Toujours Intelligent

L'équipe a testé son modèle mis à jour, qu'elle a nommé LLaDA-HYBRID, pour voir s'il était toujours aussi bon pour écrire et résoudre des problèmes, et s'il était réellement plus rapide.

Sait-il toujours comment réfléchir ?
Étonnamment, oui. Même s'ils ont remplacé 6 des 20 couches par cette méthode de « fiche de synthèse », les performances du modèle sont restées très proches de l'original.

• Sur un test de codage appelé HumanEval, le modèle original a obtenu un score de 75,6 %, et le modèle hybride de 72,0 %.
• Sur un test de mathématiques appelé CMATH, l'original a obtenu 88,3 %, et l'hybride 86,7 %.
• De manière intéressante, sur un autre test de codage, MBPP+, le modèle hybride s'est amélioré, passant de 57,7 % à 63,0 %.
  Le papier suggère que, bien que le modèle soit légèrement moins parfait sur les tâches de raisonnement les plus difficiles (comme un quiz scientifique complexe appelé GPQA-Diamond, où il est passé de 38,9 % à 30,8 %), il a largement préservé sa capacité à écrire du code et à résoudre des problèmes mathématiques.

Est-il réellement plus rapide ?
C'est ici que la magie opère. Parce que le modèle n'a plus besoin de relire l'intégralité de l'historique de la conversation pour chaque nouveau mot, il peut gérer plus de requêtes à la fois.

• Lors des tests sur le nombre de mots que le modèle pouvait générer par seconde, le modèle hybride était 1,7 fois plus rapide que l'original lors de ses pics de performance (en gérant 128 requêtes simultanées).
• L'utilisation de la mémoire du modèle original augmentait à mesure que l'histoire s'allongeait, finissant par heurter un mur où il ne pouvait plus gérer davantage d'utilisateurs. Le modèle hybride, grâce à sa « fiche de synthèse » de taille fixe, a pu gérer plus de requêtes simultanées avant de manquer de mémoire.

Ce qu'il faut retenir

Ce papier suggère que nous n'avons pas toujours besoin de construire une nouvelle IA plus rapide à partir de zéro. Au lieu de cela, nous pouvons prendre une IA puissante existante et lui donner un système de « mémoire hybride ». En gardant une concentration aiguisée sur le moment présent tout en résumant le passé sous une forme compacte, nous pouvons rendre ces modèles nettement plus rapides et plus efficaces. Les chercheurs ont constaté que cette mise à niveau pouvait être réalisée en environ 60 heures d'entraînement sur des données publiques, offrant une voie prometteuse pour rendre les services d'IA plus rapides et moins coûteux sans trop sacrifier leur intelligence. Cependant, les auteurs notent que ce n'est que le début ; ils n'ont mis à niveau que 6 couches, et les travaux futurs pourraient trouver des moyens de mettre à niveau encore plus de couches ou de gérer des histoires encore plus longues.