MoM: Linear Sequence Modeling with Mixture-of-Memories

TL;DR

MoM introduces a multi-memory architecture for linear sequence models, significantly improving recall performance while maintaining linear training complexity and constant inference complexity, bridging the gap with Transformer models.

Contribution

MoM proposes a novel multi-memory framework with a routing mechanism, enhancing memory capacity and recall ability in linear sequence models, a significant advancement over existing single-memory approaches.

Findings

MoM outperforms existing linear models on recall-intensive language tasks.

MoM achieves performance comparable to Transformer models.

MoM maintains linear training and constant inference complexity.

Abstract

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 tasks. To address this limitation, 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. MoM serves as a general framework that can be seamlessly combined with diverse memory update mechanisms across linear models. 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 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.