MoM：基于混合记忆的线性序列建模

线性注意力、状态空间建模和线性 RNN 等线性序列建模方法通过降低训练和推理的复杂度，提供了显著的效率提升。然而，这些方法通常将整个输入序列压缩为单个固定大小的记忆状态，这导致其在需要密集召回的任务上表现欠佳。为解决这一局限性，我们引入了一种名为混合记忆（Mixture-of-Memories, MoM）的新架构。MoM 利用多个独立的记忆状态，并通过路由器网络将输入 token 引导至特定的记忆状态。该方法在最小化记忆干扰的同时，极大地提升了整体记忆容量。MoM 作为一个通用框架，可以与线性模型中多种记忆更新机制无缝结合。因此，MoM 在密集召回任务上表现优异，超越了现有的线性序列建模技术。尽管结合了多个记忆状态，每个记忆状态的计算复杂度仍保持线性，使得 MoM 在训练时保持线性复杂度优势，而在推理时保持常数复杂度。我们的实验结果表明，MoM 在下游语言任务上优于当前的线性序列模型，尤其是密集召回任务，甚至取得了与 Transformer 模型相当的性能。代码发布于 https://github.com/OpenSparseLLMs/MoM，并作为 https://github.com/OpenSparseLLMs/Linear-MoE 的一部分发布。