MoM: Modelagem de Sequências Lineares com Mistura de Memórias

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

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

cs.AI

Resumo

Métodos de modelagem de sequências lineares, como atenção linear, modelagem de espaço de estados e RNNs lineares, oferecem melhorias significativas de eficiência ao reduzir a complexidade do treinamento e da inferência. No entanto, esses métodos geralmente comprimem toda a sequência de entrada em um único estado de memória de tamanho fixo, o que leva a um desempenho subótimo em tarefas subsequentes intensivas em recuperação de memória. Inspirados pela neurociência, particularmente pela capacidade do cérebro de manter memória de longo prazo robusta enquanto mitiga a "interferência de memória", introduzimos uma nova arquitetura chamada Mixture-of-Memories (MoM). O MoM utiliza múltiplos estados de memória independentes, com uma rede roteadora direcionando os tokens de entrada para estados de memória específicos. Essa abordagem aumenta consideravelmente a capacidade geral de memória enquanto minimiza a interferência de memória. Como resultado, o MoM se sai excepcionalmente bem em tarefas intensivas em recuperação, superando as técnicas existentes de modelagem de sequências lineares. Apesar de incorporar múltiplos estados de memória, o cálculo de cada estado de memória permanece linear em complexidade, permitindo que o MoM mantenha a vantagem de complexidade linear durante o treinamento, enquanto mantém complexidade constante durante a inferência. Nossos resultados experimentais mostram que o MoM supera significativamente os modelos de sequências lineares atuais em tarefas de linguagem subsequentes, particularmente em tarefas intensivas em recuperação, e até alcança desempenho comparável aos modelos Transformer. O código foi liberado em https://github.com/OpenSparseLLMs/MoM e também faz 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.