TransMamba: TransformerとMambaの柔軟な切り替え

TransMamba: Flexibly Switching between Transformer and Mamba

March 31, 2025

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

cs.AI

要旨

Transformerは現代の大規模言語モデルの基盤となっていますが、その二次的な計算複雑性が長系列処理における効率性を制限しています。最近、線形複雑性を持つ状態空間モデル（SSM）であるMambaの進展により、効率性の向上が期待されていますが、不安定な文脈学習やマルチタスク汎化の問題を抱えています。本論文では、TransformerとMambaを共有パラメータ行列（例：QKVおよびCBx）を通じて統合し、異なるトークン長や層において動的に注意機構とSSM機構を切り替えることができる新しいフレームワーク、TransMambaを提案します。我々は、TransformerとMambaを橋渡しするために、注意出力をSSM互換の状態に変換するMemory converterを設計し、変換が行われるTransPointsにおいてシームレスな情報の流れを確保します。さらに、TransPointのスケジューリングについても詳細に検討し、さらなる改善を図ります。広範な実験を通じて、TransMambaがベースラインと比較して優れた学習効率と性能を達成することを実証し、TransformerとMambaのパラダイム間の深い一貫性を検証しました。これにより、次世代の系列モデリングに向けたスケーラブルなソリューションを提供します。

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.