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 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 this https URL and is also released as a part of this https URL.
Abstract (translated)
线性序列建模方法，如线性注意、状态空间模型和线性RNNs，通过减少训练和推理的复杂度提供了显著的效率提升。然而，这些方法通常会将整个输入序列压缩到单一固定大小的记忆状态中，这在需要大量回忆能力的任务上会导致次优性能。受到神经科学的启发，特别是大脑能够在保持强大的长期记忆的同时减轻“记忆干扰”的能力，我们引入了一种新的架构叫做混合内存（MoM）。MoM利用多个独立的记忆状态，并且一个路由器网络负责将输入令牌导向特定的记忆状态。这种做法极大地提高了整体记忆力容量，同时最小化了记忆之间的干扰。因此，MoM在需要大量回忆任务上表现出色，超越现有的线性序列建模技术。 尽管采用了多内存状态，但每个内存状态的计算复杂度仍然保持为线性，在训练时保留了线性复杂度的优势，并且在推理中维持恒定复杂度。我们的实验结果表明，MoM在线下语言任务特别是需要大量回忆的任务上显著优于当前的线性序列模型，并且甚至实现了与Transformer模型相当的表现。 代码可在以下链接找到：[请在此处插入具体URL]，并且也是这个项目的组成部分之一：[请在此处插入具体URL]。
URL
https://arxiv.org/abs/2502.13685
