Meta a récemment publié les résultats d'une recherche révolutionnaire : le développement d'une nouvelle technologie de couche mémoire qui améliore considérablement la précision factuelle des grands modèles linguistiques (LLM) et permet une extension sans précédent de l'échelle des paramètres. Cette technologie remet en question les méthodes d'extension traditionnelles des réseaux neuronaux et ouvre de nouvelles perspectives pour la conception future des architectures d'IA.
Au cœur de cette recherche se trouve l'utilisation d'un mécanisme de recherche clé-valeur entraînable pour ajouter des paramètres supplémentaires au modèle sans augmenter la charge de calcul (FLOPs). L'idée principale est de compléter les couches de transmission directe gourmandes en calcul par une couche mémoire à activation clairsemée, offrant ainsi une capacité de stockage et de récupération d'informations spécialisée.
Comparée aux réseaux denses traditionnels, la couche mémoire est plus efficace pour le stockage d'informations. Par exemple, pour apprendre des informations simples comme les dates d'anniversaire ou les capitales, la couche mémoire utilise un mécanisme de recherche clé-valeur simple, ce qui est plus efficace que l'utilisation d'un réseau de transmission directe.
La principale contribution de cette recherche réside dans l'extension de l'échelle de la couche mémoire à un niveau sans précédent, atteignant 128 milliards de paramètres. Les résultats expérimentaux montrent que, dans les tâches en aval, les modèles linguistiques équipés de la couche mémoire améliorée surpassent non seulement les modèles denses dont la charge de calcul est doublée, mais aussi les modèles d'experts mixtes lorsque la charge de calcul et le nombre de paramètres sont identiques. L'amélioration des performances est particulièrement remarquable pour les tâches factuelles.
Les chercheurs de Meta ont atteint cet objectif en remplaçant un ou plusieurs réseaux de transmission directe (FFN) du réseau Transformer par une couche mémoire. Ce remplacement a montré des avantages constants pour différentes tailles de modèles de base (de 134 millions à 8 milliards de paramètres) et capacités de mémoire (jusqu'à 128 milliards de paramètres). Les résultats expérimentaux montrent que la couche mémoire peut améliorer la précision factuelle des modèles linguistiques de plus de 100 %, avec des améliorations significatives également pour la programmation et les connaissances générales. Dans de nombreux cas, les modèles équipés d'une couche mémoire atteignent même les performances de modèles denses nécessitant quatre fois plus de calculs.
Les chercheurs ont également apporté plusieurs améliorations à la couche mémoire pour surmonter les défis liés à son application à grande échelle :
Adoption d'un mécanisme de recherche clé multiplicative : pour résoudre le goulot d'étranglement de la recherche de clés de requête dans les couches mémoire à grande échelle, la recherche a utilisé des clés de quantification multiplicative entraînables, évitant ainsi la comparaison paire par paire de chaque clé de requête.
Couche mémoire parallélisée : pour permettre la parallélisation de la couche mémoire dans un environnement multi-GPU, les chercheurs ont réparti les opérations de recherche et d'agrégation d'intégration sur plusieurs GPU.
Mécanisme de mémoire partagée : pour maximiser le partage des paramètres, les chercheurs utilisent un pool de paramètres de mémoire partagé entre toutes les couches mémoire.
Optimisation des performances et de la stabilité : les chercheurs ont optimisé l'opération EmbeddingBag à l'aide de noyaux CUDA personnalisés, améliorant considérablement l'utilisation de la bande passante mémoire. Ils ont également introduit un mécanisme de contrôle d'entrée lié à l'entrée avec une non-linéarité silu pour améliorer les performances et la stabilité de l'entraînement.
Les résultats expérimentaux ont également révélé les découvertes clés suivantes :
La taille de la couche mémoire a un impact significatif sur les performances : plus la couche mémoire est grande, plus les performances de la réponse aux questions factuelles sont élevées.
Plusieurs couches mémoire sont supérieures à une seule couche mémoire : l'utilisation de plusieurs couches mémoire avec des paramètres partagés peut améliorer les performances, mais un nombre excessif de couches mémoire peut les diminuer. Le nombre optimal de couches mémoire est de trois.
La couche mémoire apprend les faits plus rapidement : au début de l'entraînement, les modèles équipés d'une couche mémoire présentent une amélioration des performances plus rapide, ce qui indique que la couche mémoire aide le modèle à apprendre les faits plus rapidement.
La couche mémoire et la couche dense sont complémentaires : les expériences montrent que les couches mémoire clairsemées et les couches de transmission directe denses sont toutes deux essentielles.
Pour vérifier l'efficacité de la technologie de couche mémoire, les chercheurs ont effectué des évaluations sur plusieurs tests de référence, notamment :
Réponse aux questions factuelles (NaturalQuestions, TriviaQA)
Réponse aux questions à sauts multiples (HotpotQA)
Connaissances scientifiques et générales (MMLU, HellaSwag, OBQA, PIQA)
Programmation (HumanEval, MBPP)
Les résultats montrent que les modèles équipés d'une couche mémoire surpassent les modèles de référence dans ces tests, l'amélioration étant particulièrement marquée pour la réponse aux questions factuelles.
Cette recherche de Meta non seulement offre de nouvelles perspectives pour l'extension des modèles d'IA, mais ouvre également de nouvelles voies pour résoudre les problèmes factuels et améliorer les performances des modèles. Les chercheurs estiment que la technologie de couche mémoire est hautement évolutive et devrait être largement utilisée dans diverses applications d'IA à l'avenir. Ils soulignent également que, bien que la couche mémoire présente encore des défis en termes d'accélération matérielle, ils sont convaincus que des recherches et des optimisations continues permettront d'atteindre, voire de dépasser, les performances des réseaux de transmission directe traditionnels.
De plus, l'équipe de recherche de Meta espère améliorer les performances de la couche mémoire grâce à de nouvelles méthodes d'apprentissage, réduire l'oubli et les hallucinations du modèle et permettre l'apprentissage continu.
La publication de cette recherche a sans aucun doute injecté une nouvelle vitalité dans le domaine de l'IA, et nous remplit d'espoir quant à l'avenir de l'IA.
Article : https://arxiv.org/pdf/2412.09764
