TransMamba: Flexibly Switching between Transformer and Mamba
This work addresses a key bottleneck in sequence modeling for AI applications by offering a scalable solution that improves efficiency and stability, though it appears incremental as it builds on existing Transformer and Mamba paradigms.
The paper tackles the inefficiency of Transformers in long-sequence processing and the instability of Mamba models by proposing TransMamba, a framework that dynamically switches between Transformer and Mamba mechanisms, achieving superior training efficiency and performance compared to baselines.
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.