Nouvelle architecture LLM de Google : séparation des mémoires pour maîtriser coûts et performances

Une nouvelle architecture de réseau de neurones développée par des chercheurs de Google pourrait résoudre un des grands défis rencontrés par les modèles de langage de grande taille (LLMs) : étendre leur mémoire lors de l’inférence sans faire exploser les coûts associés à la mémoire et aux calculs. Nommée Titans, cette architecture permet aux modèles de trouver et de stocker, durant l’inférence, de petites informations cruciales au sein de longues séquences.

Titans combine des blocs d’attention traditionnels de LLM avec des couches de « mémoire neurale », permettant aux modèles de gérer efficacement des tâches de mémoire à court et à long terme. Selon les chercheurs, les LLMs utilisant la mémoire neurale à long terme peuvent se scaler jusqu’à des millions de tokens tout en surpassant les LLM classiques et des alternatives comme Mamba, avec un nombre de paramètres réduit.

Couches d’attention et modèles linéaires

Les architectures classiques de transformateurs utilisées dans les LLMs emploient le mécanisme de self-attention pour établir les relations entre les tokens. C’est une technique efficace pour apprendre des motifs complexes dans les séquences de tokens. Cependant, à mesure que la longueur de la séquence augmente, les coûts en mémoire et en calcul pour le calcul et le stockage de l’attention s’accroissent de manière quadratique.

Des propositions récentes ont présenté des architectures alternatives avec une complexité linéaire, capables de s’adapter sans gonfler les coûts de mémoire et de calcul. Toutefois, les chercheurs de Google soutiennent que les modèles linéaires ne parviennent pas à rivaliser avec les transformateurs classiques en termes de performance, car ils compriment leurs données contextuelles et peuvent manquer des détails importants.

Ils suggèrent qu’une architecture idéale devrait comporter différents composants mémoriels pouvant interagir pour utiliser les connaissances existantes, mémoriser de nouveaux faits et apprendre des abstractions à partir de leur contexte.

« Nous soutenons que, dans un paradigme d’apprentissage efficace, à l’instar du cerveau humain, il existe des modules distincts mais interconnectés, chacun étant responsable d’un composant essentiel au processus d’apprentissage », écrivent les chercheurs.

Mémoire neurale à long terme

« La mémoire est une organisation de systèmes — par exemple, mémoire à court terme, mémoire de travail, et mémoire à long terme — chacun jouant une fonction différente avec des structures neuronales variées, et chacun capable d’opérer de manière indépendante », notent les chercheurs.

Pour combler les lacunes des modèles de langue actuels, les chercheurs proposent un module de « mémoire neurale à long terme » capable d’apprendre de nouvelles informations durant l’inférence sans les inefficacités du mécanisme d’attention complet. Au lieu de stocker des informations pendant l’entraînement, le module de mémoire neurale apprend une fonction capable de mémoriser de nouveaux faits pendant l’inférence et d’adapter dynamiquement le processus de mémorisation en fonction des données rencontrées.

Pour décider quelles informations méritent d’être stockées, le module de mémoire neurale utilise le concept de « surprise ». Plus une séquence de tokens diffère des informations déjà intégrées dans les poids et la mémoire existante du modèle, plus elle est surprenante et donc digne d’être mémorisée. Cela permet au module d’utiliser efficacement sa mémoire limitée et de ne conserver que les données ajoutant une information utile à ses connaissances préalables.

Pour gérer de très longues séquences de données, le module de mémoire neurale possède un mécanisme d’oubli adaptatif qui lui permet d’éliminer l’information devenue obsolète, contribuant ainsi à la gestion de la capacité limitée de la mémoire.

Ce module mémoriel peut compléter le mécanisme d’attention des modèles transformateurs actuels, que les chercheurs décrivent comme des « modules de mémoire à court terme, prenant en compte la taille de la fenêtre contextuelle actuelle. En revanche, notre mémoire neurale, capable d’apprendre en continu à partir des données et de les stocker dans ses poids, peut jouer le rôle d’une mémoire à long terme ».

Architecture Titan

Les chercheurs décrivent Titans comme une famille de modèles intégrant des blocs transformateurs existants avec des modules de mémoire neurale. Le modèle se compose de trois éléments principaux : le module « central », qui agit comme une mémoire à court terme et utilise le mécanisme d’attention classique pour traiter le segment actuel des tokens d’entrée ; un module de « mémoire à long terme », exploitant l’architecture de mémoire neurale pour conserver des informations au-delà du contexte immédiat ; et un module de « mémoire persistante », qui représente les paramètres apprenables fixes après l’entraînement, stockant des connaissances non sensibles au temps.

Les chercheurs proposent différentes manières de connecter ces trois composants. L’avantage principal de cette architecture réside dans la synergie entre les modules d’attention et de mémoire. Par exemple, les couches d’attention peuvent utiliser les contextes historique et actuel pour déterminer quais aspects de la fenêtre contextuelle actuelle doivent être enregistrés dans la mémoire à long terme. Pendant ce temps, la mémoire à long terme fournit des connaissances historiques qui ne figurent pas dans le contexte actuel d’attention.

Les chercheurs ont réalisé des tests à petite échelle sur les modèles Titan, allant de 170 millions à 760 millions de paramètres, sur diverses tâches, notamment la modélisation linguistique et les tâches de langage en longues séquences. Ils ont comparé les performances des Titans à celles de divers modèles basés sur des transformateurs, de modèles linéaires tels que Mamba et de modèles hybrides tels que Samba.

Les Titans ont montré de bonnes performances en modélisation linguistique par rapport à d’autres modèles, surpassant les transformateurs ainsi que les modèles linéaires de taille comparable.

La différence de performance est particulièrement marquée dans des tâches exigeant de longues séquences, comme celle de « l’aiguille dans une botte de foin », où le modèle doit extraire des informations d’une très longue séquence, et BABILong, où il doit raisonner à partir de faits dispersés dans des documents volumiques. En effet, sur ces tâches, Titan a dépassé des modèles disposant de paramètres largement supérieurs, comme GPT-4 et GPT-4o-mini, ainsi qu’un modèle Llama-3 amélioré par génération augmentée par récupération (RAG).

De plus, les chercheurs ont pu étendre la fenêtre contextuelle des Titans jusqu’à 2 millions de tokens tout en maintenant les coûts de mémoire à un niveau raisonnable.

Si les modèles doivent encore être testés sur des tailles plus importantes, les résultats obtenus montrent que les chercheurs n’ont pas encore atteint le plein potentiel des Titans.

Quelle implication pour les applications d’entreprise ?

Étant donné que Google se positionne à la pointe des modèles à long contexte, cela laisse supposer que cette technique fera probablement son apparition dans des modèles privés et ouverts tels que Gemini et Gemma.

Avec des LLMs capables de traiter de plus longues fenêtres contextuelles, le potentiel de création d’applications où l’on peut condenser des connaissances nouvelles dans une invite plutôt que d’utiliser des techniques comme RAG est en pleine expansion. Le cycle de développement pour créer et itérer des applications basées sur des invites est bien plus rapide comparé aux pipelines RAG complexes. Par ailleurs, des architectures comme Titans peuvent contribuer à réduire les coûts d’inférence pour des séquences très longues, rendant possible le déploiement d’applications LLM par les entreprises dans un plus grand nombre de cas d’utilisation.

Google prévoit de publier le code PyTorch et JAX pour l’entraînement et l’évaluation des modèles Titans.