arXivDaily arXiv每日学术速递 周一至周五更新
​通知:联邦节假日调休:哥伦布日 / 原住民日(Columbus Day / Indigenous Peoples' Day),今日(13日)无更新。
arXiv 2502.13685cs.CLcs.AIcs.LG

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

MoM: Linear Sequence Modeling with Mixture-of-Memories

Shanghai AI Laboratory(上海人工智能实验室);Tsinghua University(清华大学);Fudan University(复旦大学);The Hong Kong University of Science and Technology (Guangzhou)(香港科技大学(广州));The Chinese University of Hong Kong(香港中文大学)

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

AI总结:

针对现有线性序列建模方法单固定记忆状态导致强召回任务表现差的问题,提出MoM混合记忆架构,通过路由分配多独立记忆状态提升容量、降低干扰,保持线性训练与常数推理复杂度,性能优于现有线性模型且比肩Transformer。

摘要:

线性注意力、状态空间建模、线性循环神经网络(RNN)等线性序列建模方法通过降低训练与推理复杂度,实现了显著的效率提升。但这类方法通常会将整个输入序列压缩为单个固定大小的记忆状态,导致其在强记忆召回类任务上表现欠佳。为解决这一缺陷,我们提出了名为混合记忆(Mixture-of-Memories,MoM)的新型架构。MoM采用多个独立记忆状态,通过路由网络将输入token分配至对应的特定记忆状态,该方式在大幅提升整体记忆容量的同时可最小化记忆干扰。MoM是一种通用框架,能够与线性模型中各类记忆更新机制无缝结合,因此在强记忆召回类任务上表现优异,性能超越现有线性序列建模技术。尽管引入了多个记忆状态,每个记忆状态的计算仍保持线性复杂度,使MoM在训练阶段保留线性复杂度优势,推理阶段则具备常数复杂度。实验结果显示,MoM在下游语言任务尤其是强记忆召回类任务上性能优于当前线性序列模型,甚至可达到与Transformer模型相当的表现。代码已发布于https://github.com/OpenSparseLLMs/MoM,同时也作为https://github.com/OpenSparseLLMs/Linear-MoE的一部分公开。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑