TransMamba: Flexibles Wechseln zwischen Transformer und Mamba

TransMamba: Flexibly Switching between Transformer and Mamba

March 31, 2025

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

cs.AI

Zusammenfassung

Transformer sind der Eckpfeiler moderner großer Sprachmodelle, doch ihre quadratische Rechenkomplexität schränkt die Effizienz bei der Verarbeitung langer Sequenzen ein. Jüngste Fortschritte bei Mamba, einem Zustandsraummodell (State Space Model, SSM) mit linearer Komplexität, bieten vielversprechende Effizienzgewinne, leiden jedoch unter instabilem kontextuellem Lernen und mangelnder Generalisierung bei Mehrfachaufgaben. In diesem Artikel wird TransMamba vorgeschlagen, ein neuartiges Framework, das Transformer und Mamba durch gemeinsame Parametermatrizen (z. B. QKV und CBx) vereint und somit dynamisch zwischen Aufmerksamkeits- und SSM-Mechanismen bei unterschiedlichen Token-Längen und Schichten wechseln kann. Wir entwickeln den Memory Converter, um Transformer und Mamba zu verbinden, indem Aufmerksamkeitsausgaben in SSM-kompatible Zustände umgewandelt werden, wodurch ein nahtloser Informationsfluss an den TransPoints, an denen die Transformation stattfindet, sichergestellt wird. Die TransPoint-Planung wird ebenfalls umfassend untersucht, um weitere Verbesserungen zu erzielen. Wir führten umfangreiche Experimente durch, die zeigen, dass TransMamba eine überlegene Trainings effizienz und Leistung im Vergleich zu den Baselines erreicht, und bestätigten die tiefere Konsistenz zwischen den Transformer- und Mamba-Paradigmen, was eine skalierbare Lösung für die nächste Generation der Sequenzmodellierung bietet.

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.