TransMamba: Flexibel schakelen tussen Transformer en Mamba

TransMamba: Flexibly Switching between Transformer and Mamba

March 31, 2025

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

cs.AI

Samenvatting

Transformers vormen de hoeksteen van moderne grote taalmodellen, maar hun kwadratische computationele complexiteit beperkt de efficiëntie bij de verwerking van lange sequenties. Recente vooruitgang in Mamba, een state space model (SSM) met lineaire complexiteit, biedt veelbelovende efficiëntiewinsten, maar lijdt onder instabiel contextueel leren en multitask-generalizatie. Dit artikel stelt TransMamba voor, een nieuw raamwerk dat Transformer en Mamba verenigt via gedeelde parameter matrices (bijv. QKV en CBx), en daardoor dynamisch kan schakelen tussen aandacht- en SSM-mechanismen bij verschillende tokenlengtes en lagen. We ontwerpen de Memory converter om Transformer en Mamba te verbinden door aandachtsuitvoer om te zetten in SSM-compatibele toestanden, wat een naadloze informatieflow garandeert bij TransPoints waar de transformatie plaatsvindt. De TransPoint-scheduling wordt ook grondig onderzocht voor verdere verbeteringen. We hebben uitgebreide experimenten uitgevoerd die aantonen dat TransMamba superieure trainings efficiëntie en prestaties bereikt in vergelijking met baselines, en hebben de diepere consistentie tussen de Transformer- en Mamba-paradigma's gevalideerd, wat een schaalbare oplossing biedt voor next-generation sequentiemodellering.

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. 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.