MoM : Modélisation de séquences linéaires avec un mélange de mémoires

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

Auteurs: Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Yu Cheng

cs.AI

Résumé

Les méthodes de modélisation de séquences linéaires, telles que l'attention linéaire, la modélisation par espace d'états et les RNN linéaires, offrent des améliorations significatives en termes d'efficacité en réduisant la complexité de l'entraînement et de l'inférence. Cependant, ces méthodes compressent généralement l'intégralité de la séquence d'entrée en un seul état de mémoire de taille fixe, ce qui entraîne des performances sous-optimales pour les tâches en aval intensives en rappel. S'inspirant des neurosciences, en particulier de la capacité du cerveau à maintenir une mémoire à long terme robuste tout en atténuant les "interférences de mémoire", nous introduisons une nouvelle architecture appelée Mixture-of-Memories (MoM). MoM utilise plusieurs états de mémoire indépendants, avec un réseau de routage qui dirige les tokens d'entrée vers des états de mémoire spécifiques. Cette approche améliore considérablement la capacité de mémoire globale tout en minimisant les interférences de mémoire. En conséquence, MoM excelle dans les tâches intensives en rappel, surpassant les techniques existantes de modélisation de séquences linéaires. Bien qu'elle intègre plusieurs états de mémoire, le calcul de chaque état de mémoire reste de complexité linéaire, permettant à MoM de conserver l'avantage de la complexité linéaire pendant l'entraînement, tout en maintenant une complexité constante pendant l'inférence. Nos résultats expérimentaux montrent que MoM surpasse significativement les modèles de séquences linéaires actuels sur les tâches linguistiques en aval, en particulier les tâches intensives en rappel, et atteint même des performances comparables à celles des modèles Transformer. Le code est disponible à l'adresse https://github.com/OpenSparseLLMs/MoM et fait également partie de https://github.com/OpenSparseLLMs/Linear-MoE.

English

Linear sequence modeling methods, such as linear attention, state space modeling, and linear RNNs, offer significant efficiency improvements by reducing the complexity of training and inference. However, these methods typically compress the entire input sequence into a single fixed-size memory state, which leads to suboptimal performance on recall-intensive downstream tasks. Drawing inspiration from neuroscience, particularly the brain's ability to maintain robust long-term memory while mitigating "memory interference", we introduce a novel architecture called Mixture-of-Memories (MoM). MoM utilizes multiple independent memory states, with a router network directing input tokens to specific memory states. This approach greatly enhances the overall memory capacity while minimizing memory interference. As a result, MoM performs exceptionally well on recall-intensive tasks, surpassing existing linear sequence modeling techniques. Despite incorporating multiple memory states, the computation of each memory state remains linear in complexity, allowing MoM to retain the linear-complexity advantage during training, while constant-complexity during inference. Our experimental results show that MoM significantly outperforms current linear sequence models on downstream language tasks, particularly recall-intensive tasks, and even achieves performance comparable to Transformer models. The code is released at https://github.com/OpenSparseLLMs/MoM and is also released as a part of https://github.com/OpenSparseLLMs/Linear-MoE.