TransMamba : Commutation flexible entre Transformer et 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

Résumé

Les Transformers constituent la pierre angulaire des grands modèles de langage modernes, mais leur complexité computationnelle quadratique limite leur efficacité dans le traitement de séquences longues. Les avancées récentes de Mamba, un modèle d'espace d'état (SSM) à complexité linéaire, offrent des gains d'efficacité prometteurs mais souffrent d'un apprentissage contextuel instable et d'une généralisation multitâche limitée. Cet article propose TransMamba, un cadre novateur qui unifie Transformer et Mamba à travers des matrices de paramètres partagées (par exemple, QKV et CBx), permettant ainsi de basculer dynamiquement entre les mécanismes d'attention et de SSM à différentes longueurs de tokens et couches. Nous concevons le Convertisseur de Mémoire pour relier Transformer et Mamba en convertissant les sorties d'attention en états compatibles avec SSM, assurant un flux d'information fluide aux TransPoints où la transformation s'opère. La planification des TransPoints est également explorée en profondeur pour des améliorations supplémentaires. Nous avons mené des expériences approfondies démontrant que TransMamba atteint une efficacité d'entraînement et des performances supérieures par rapport aux modèles de référence, et validé la cohérence profonde entre les paradigmes Transformer et Mamba, offrant ainsi une solution évolutive pour la modélisation de séquences de nouvelle génération.

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.