MoM: Modellazione di Sequenze Lineari con Miscela di Memorie

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

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

cs.AI

Abstract

I metodi di modellazione di sequenze lineari, come l'attenzione lineare, la modellazione dello spazio degli stati e le RNN lineari, offrono significativi miglioramenti in termini di efficienza riducendo la complessità dell'addestramento e dell'inferenza. Tuttavia, questi metodi comprimono tipicamente l'intera sequenza di input in un unico stato di memoria di dimensione fissa, il che porta a prestazioni subottimali nei task downstream intensivi di richiamo. Traendo ispirazione dalle neuroscienze, in particolare dalla capacità del cervello di mantenere una memoria a lungo termine robusta mitigando l'"interferenza di memoria", introduciamo una nuova architettura chiamata Mixture-of-Memories (MoM). MoM utilizza più stati di memoria indipendenti, con una rete di routing che indirizza i token di input a specifici stati di memoria. Questo approccio migliora notevolmente la capacità complessiva della memoria minimizzando l'interferenza di memoria. Di conseguenza, MoM si comporta eccezionalmente bene nei task intensivi di richiamo, superando le tecniche esistenti di modellazione di sequenze lineari. Nonostante l'incorporazione di più stati di memoria, il calcolo di ciascuno stato di memoria rimane lineare in complessità, consentendo a MoM di mantenere il vantaggio della complessità lineare durante l'addestramento e una complessità costante durante l'inferenza. I nostri risultati sperimentali mostrano che MoM supera significativamente i modelli di sequenze lineari attuali nei task linguistici downstream, in particolare quelli intensivi di richiamo, e raggiunge persino prestazioni comparabili ai modelli Transformer. Il codice è rilasciato su https://github.com/OpenSparseLLMs/MoM e fa anche parte di 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.