提出可动态切换的混合架构,让大模型处理长文本更高效。
TransMamba: A Sequence-Level Hybrid Transformer-Mamba Language Model
- 用共享参数在序列级融合Transformer与Mamba
- 训练效率比基线高30%,长序列性能提升显著
- 适合追求长文本处理效率的研究者和开发者
Transformer是现代大语言模型的核心,但其二次计算复杂度限制了长序列处理效率。近期基于状态空间模型(SSM)的Mamba以线性复杂度展现高效潜力,却存在上下文学习不稳定和多任务泛化差的问题。部分工作尝试在层级别混合Transformer与Mamba结构,以兼顾两者优势。本文提出TransMamba,一种新的序列级混合框架,通过共享参数矩阵(QKV与CBx)统一Transformer与Mamba,实现不同词元长度与层间动态切换注意力与SSM机制。设计了记忆转换器(Memory Converter),将注意力输出转化为SSM兼容状态,保障在转换点(TransPoint)的信息连续流动。同时系统探索了TransPoint调度策略,在效果与效率间取得平衡。大量实验表明,TransMamba在训练效率与性能上均优于单模型及混合基线,验证了Transformer与Mamba在序列层面的深层一致性,为下一代语言建模提供了可扩展方案。代码与数据已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。