MoM: Lineaire Sequentiële Modellering met Mixture-of-Memories

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

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

cs.AI

Samenvatting

Lineaire sequentiemodelleringsmethoden, zoals lineaire aandacht, toestandsruimtemodellering en lineaire RNN's, bieden aanzienlijke efficiëntieverbeteringen door de complexiteit van training en inferentie te verminderen. Deze methoden comprimeren echter doorgaans de volledige invoerreeks in een enkele vaste geheugenstatus, wat leidt tot suboptimale prestaties bij recall-intensieve downstreamtaken. Geïnspireerd door neurowetenschappen, met name het vermogen van de hersenen om robuust langetermijngeheugen te behouden terwijl "geheugeninterferentie" wordt geminimaliseerd, introduceren we een nieuwe architectuur genaamd Mixture-of-Memories (MoM). MoM maakt gebruik van meerdere onafhankelijke geheugenstatussen, waarbij een routeringsnetwerk invoertokens naar specifieke geheugenstatussen leidt. Deze aanpak verbetert de algehele geheugencapaciteit aanzienlijk en minimaliseert tegelijkertijd geheugeninterferentie. Als gevolg presteert MoM uitstekend bij recall-intensieve taken en overtreft het bestaande lineaire sequentiemodelleringsmethoden. Ondanks het gebruik van meerdere geheugenstatussen blijft de berekening van elke geheugenstatus lineair in complexiteit, waardoor MoM het voordeel van lineaire complexiteit tijdens de training behoudt, terwijl het tijdens de inferentie constant in complexiteit blijft. Onze experimentele resultaten tonen aan dat MoM huidige lineaire sequentiemodellen aanzienlijk overtreft bij downstreamtaaltaken, met name recall-intensieve taken, en zelfs prestaties bereikt die vergelijkbaar zijn met Transformer-modellen. De code is vrijgegeven op https://github.com/OpenSparseLLMs/MoM en maakt ook deel uit van 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.