MoM: Modelado de Secuencias Lineales con Mezcla de Memorias

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

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

cs.AI

Resumen

Los métodos de modelado de secuencias lineales, como la atención lineal, el modelado de espacios de estado y las RNNs lineales, ofrecen mejoras significativas en eficiencia al reducir la complejidad del entrenamiento y la inferencia. Sin embargo, estos métodos suelen comprimir toda la secuencia de entrada en un único estado de memoria de tamaño fijo, lo que resulta en un rendimiento subóptimo en tareas posteriores intensivas en recuperación de información. Inspirándonos en la neurociencia, particularmente en la capacidad del cerebro para mantener una memoria a largo plazo robusta mientras mitiga la "interferencia de memoria", introducimos una nueva arquitectura llamada Mezcla de Memorias (Mixture-of-Memories, MoM). MoM utiliza múltiples estados de memoria independientes, con una red enrutadora que dirige los tokens de entrada a estados de memoria específicos. Este enfoque mejora considerablemente la capacidad de memoria general mientras minimiza la interferencia de memoria. Como resultado, MoM tiene un rendimiento excepcional en tareas intensivas en recuperación, superando las técnicas existentes de modelado de secuencias lineales. A pesar de incorporar múltiples estados de memoria, el cálculo de cada estado de memoria sigue siendo lineal en complejidad, lo que permite a MoM mantener la ventaja de complejidad lineal durante el entrenamiento, y de complejidad constante durante la inferencia. Nuestros resultados experimentales muestran que MoM supera significativamente a los modelos de secuencias lineales actuales en tareas de lenguaje posteriores, especialmente en tareas intensivas en recuperación, e incluso alcanza un rendimiento comparable a los modelos Transformer. El código se ha publicado en https://github.com/OpenSparseLLMs/MoM y también forma parte 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.