Google et KAIST présentent MoR : une nouvelle architecture pour modèles linguistiques qui double la vitesse d'inférence avec 50% moins de mémoire

Une nouvelle architecture pour les modèles linguistiques : le Mixture-of-Recursions (MoR) Récemment, une équipe conjointe de chercheurs de Google DeepMind et de l'Institut de Technologie de Corée (KAIST) a publié un nouveau cadre de modèle linguistique baptisé "Mixture-of-Recursions" (MoR). Cette architecture permettrait, selon les auteurs, de doubler la vitesse d'inférence, de réduire la quantité de calculs de formation nécessaires et de diminuer l'utilisation de la mémoire KV cache de près de 50 %, tout en préservant les performances du modèle. Le travail a rapidement suscité une attention considérable sur les réseaux sociaux, certains commentaires allant jusqu'à suggérer que MoR pourrait être un "Tueur de Transformer". Depuis son introduction en 2017, l'architecture Transformer est devenue un pilier essentiel des modèles linguistiques de grande envergure. Pratiquement tous les modèles avancés actuels s'appuient sur cette structure. Cependant, la croissance exponentielle des modèles a entraîné une augmentation correspondante des exigences en ressources de calcul et de mémoire, ce qui rend leur entraînement et leur déploiement de plus en plus coûteux. Les approches traditionnelles de l'optimisation de l'efficacité, comme le partage de paramètres pour réduire la taille du modèle ou le calcul adaptatif pour allouer la puissance de calcul en fonction des besoins, n'ont généralement pas réussi à optimiser plusieurs aspects d'efficacité simultanément. Le cœur de l'innovation de l'architecture MoR réside dans l'intégration de mécanismes de calcul récursif et de routage dynamique dans un seul cadre. Ce mélange permet de traiter chaque token d'un texte d'entrée avec un nombre de couches de calcul différent, en fonction de sa complexité. Contrairement aux modèles Transformer standards, où chaque token passe par le même nombre de couches, MoR adapte la profondeur de traitement en fonction des caractéristiques spécifiques de chaque token. Pour ce faire, MoR utilise des blocs de paramètres partagés pour augmenter l'efficacité paramétrique. Un composant clé, appelé "routeur", détermine combien de fois chaque token doit subir des calculs récursifs. L'équipe de recherche a testé plusieurs stratégies de routage, notamment "expert-choice" et "token-choice", afin d'équilibrer la charge de calcul et de prévenir les problèmes logiques dans le traitement de l'information. Les tests ont montré que la stratégie "Middle-Cycle" offrait les meilleurs résultats en termes de partage de paramètres. Cette méthode conserve des paramètres indépendants uniquement pour la première et la dernière couche, tout en partageant les poids entre les couches intermédiaires. Ce compromis optimise l'efficacité paramétrique tout en maintenant la capacité expressive du modèle. La gestion de la mémoire constitue une autre amélioration significative apportée par MoR. Bien que les paramètres soient partagés, les modèles récursifs conventionnels génèrent des caches KV indépendants à chaque étape récursive, ce qui entraîne un usage élevé de la mémoire. MoR propose deux nouvelles stratégies pour résoudre ce problème : Récurrence en Cache : Ce mécanisme stocke les données KV uniquement pour les tokens qui sont routés vers des étapes récursives spécifiques. Les calculs d'attention sont limités à ces données locales, réduisant ainsi considérablement l'occupation de la mémoire KV et les opérations de lecture/écriture de données. Partage Récursif : Cette approche va plus loin en profitant du fait que tous les tokens passent par le premier bloc récursif. Elle stocke les données KV une seule fois dans la première étape, les réutilisant ensuite pour toutes les étapes récursives ultérieures, ce qui permet une économie maximale en mémoire. Les performances de MoR ont été évaluées sur plusieurs tailles de modèles, allant de 135 millions à 1,7 milliard de paramètres. Les résultats montrent que, malgré un nombre de paramètres inférieur de près de moitié, les modèles MoR dépassent les modèles Transformer traditionnels en termes de précision moyenne dans des tâches de apprentissage avec peu d'exemples. Par exemple, un modèle MoR de 384 millions de paramètres a atteint une précision moyenne de 43,1 % contre 42,3 % pour le modèle de référence Transformer. De plus, l'efficacité accrue en termes de calculs permet à MoR de traiter plus de données d'entraînement dans le même budget de calcul. Dans des expériences de comparaison avec un volume de données d'entraînement fixe, un modèle configuré avec MoR a surpassé le modèle de base alors même qu'il utilisait 25 % de calculs de formation en moins, réduisant le temps de formation de 19 % et l'usage maximal de la mémoire de 25 %. En ce qui concerne les performances d'inférence, MoR se distingue encore plus. Il utilise une technique de batchage de profondeur continue, qui regroupe les tokens à différents stades de calcul dans une même batch, facilitant le traitement grâce au partage des mêmes blocs de paramètres. Couplé à un mécanisme de sortie anticipée, ce processus améliore considérablement le débit de traitement. Ainsi, pour un modèle de 384 millions de paramètres, la configuration MoR-4 a réalisé une accélération de l'inférence de 2,06 fois sous certaines conditions. L'équipe de recherche a également constaté que le modèle MoR attribue plus de récurrences—c'est-à-dire plus de calculs—toaux tokens contenant plus d'informations sémantiques, tels que "People" ou "confident défensivement", qui nécessitent trois récurrences, contrairement aux mots fonctionnels comme "and", qui ne requièrent souvent qu'une seule. Cette observation indique que le modèle apprend à allouer les ressources de calcul de manière plus efficace et ciblée, se concentrant sur les informations les plus importantes. Les fondements de l'architecture MoR remontent aux travaux précédents menés par Google DeepMind. Les technologies comme Mixture-of-Depths (MoD), qui explorent la distribution dynamique des ressources de calcul, ont posé les bases théoriques pour MoR. De même, les Transformers récurrents, qui partagent des paramètres, ont fourni une infrastructure technique indispensable. En somme, bien que les conclusions définitives sur la pertinence et la capacité de MoR à remplacer complètement les Transformers traditionnels restent prématurées, cette architecture offre un potentiel considérable pour optimiser les performances et l'efficacité des futurs modèles linguistiques. En réduisant les coûts de déploiement et d'utilisation des grands modèles linguistiques, MoR peut aider à rendre ces technologies plus accessibles et viables à grande échelle. Références : 1. https://arxiv.org/abs/2507.10524 Préparation et mise en page : He Chenlong