arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2503.24067cs.LG

TransMamba: 一种序列级混合Transformer-Mamba语言模型

TransMamba: A Sequence-Level Hybrid Transformer-Mamba Language Model

Tencent Hunyuan(腾讯文英);The Chinese University of Hong Kong(香港中文大学);University of Macau(澳门大学)

Yixing Li, Ruobing Xie, Zhen Yang, Xingwu Sun, Shuaipeng Li, Weidong Han, Zhanhui Kang, Yu Cheng, Chengzhong Xu, Di Wang, Jie Jiang

AI总结:

TransMamba通过共享参数矩阵统一Transformer和Mamba,实现序列级动态切换注意力与SSM机制,提升训练效率和性能。

摘要:

Transformer是现代大语言模型的基石,但其二次计算复杂度限制了长序列处理的效率。最近,Mamba作为一种状态空间模型(SSM)具有线性复杂度,提供了有前途的效率提升,但存在不稳定的情境学习和多任务泛化问题。一些工作在层级上混合Transformer和Mamba层,旨在充分利用两者的优势。本文提出了TransMamba,一种新颖的序列级混合框架,通过共享参数矩阵(QKV和CBx)统一Transformer和Mamba,并因此可以在不同的token长度和层中动态切换注意力和SSM机制。我们设计了Memory Converter来连接Transformer和Mamba,通过将注意力输出转换为SSM兼容的状态,确保在TransPoints处的信息流无缝进行。TransPoint调度也得到了彻底探讨,以平衡有效性和效率。我们进行了广泛的实验,证明TransMamba在训练效率和性能上优于单一和混合基线,并验证了在序列层面Transformer和Mamba范式之间更深层次的一致性,提供了一种可扩展的解决方案,用于下一代语言建模。代码和数据可在https://github.com/Yixing-Li/TransMamba获取。

英文摘要:

Transformers are the cornerstone of modern large language models, but their quadratic computational complexity limits efficiency in long-sequence processing. Recent advancements in Mamba, a state space model (SSM) with linear complexity, offer promising efficiency gains but suffer from unstable contextual learning and multitask generalization. Some works conduct layer-level hybrid structures that combine Transformer and Mamba layers, aiming to make full use of both advantages. This paper proposes TransMamba, a novel sequence-level hybrid framework that unifies Transformer and Mamba through shared parameter matrices (QKV and CBx), and thus could dynamically switch between attention and SSM mechanisms at different token lengths and layers. We design the Memory Converter to bridge Transformer and Mamba by converting attention outputs into SSM-compatible states, ensuring seamless information flow at TransPoints where the transformation happens. The TransPoint scheduling is also thoroughly explored for balancing effectiveness and efficiency. We conducted extensive experiments demonstrating that TransMamba achieves superior training efficiency and performance compared to single and hybrid baselines, and validated the deeper consistency between Transformer and Mamba paradigms at sequence level, offering a scalable solution for next-generation language modeling. Code and data are available at https://github.com/Yixing-Li/TransMamba

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

↑