TransMamba: Flexibly Switching between Transformer and Mamba
Authors: name: Yixing Li; name: Ruobing Xie; name: Zhen Yang; name: Xingwu Sun; name: Shuaipeng Li; name: Weidong Han; name: Zhanhui Kang; name: Yu Cheng; name: Chengzhong Xu; name: Di Wang; name: Jie Jiang
Published: 2025-03-31 13:26:24

Abstract: 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. This paper proposes
TransMamba, a novel framework that unifies Transformer and Mamba through shared
parameter matrices (e.g., 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 further improvements. We conducted extensive
experiments demonstrating that TransMamba achieves superior training efficiency
and performance compared to baselines, and validated the deeper consistency
between Transformer and Mamba paradigms, offering a scalable solution for
next-generation sequence modeling.