Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TransMamba: Flexibly Switching between Transformer and Mamba

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

TL;DR

TransMamba tackles the inefficiency of Transformer‑based long‑context processing by unifying Transformer and Mamba through a single shared parameter set for $QKV$ and $CBx$. The method introduces a lossless Memory Converter that transfers information at TransPoints where the switch between attention and SSM occurs, enabling seamless transitions. It provides a systematic TransPoint scheduling strategy and demonstrates that the model achieves superior training efficiency and competitive or superior performance across eight general tasks and ultra‑long benchmarks. The work reveals deeper consistency between Transformer and Mamba and offers a scalable approach for next‑generation sequence modeling.

Abstract

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.

TransMamba: Flexibly Switching between Transformer and Mamba

TL;DR

TransMamba tackles the inefficiency of Transformer‑based long‑context processing by unifying Transformer and Mamba through a single shared parameter set for and . The method introduces a lossless Memory Converter that transfers information at TransPoints where the switch between attention and SSM occurs, enabling seamless transitions. It provides a systematic TransPoint scheduling strategy and demonstrates that the model achieves superior training efficiency and competitive or superior performance across eight general tasks and ultra‑long benchmarks. The work reveals deeper consistency between Transformer and Mamba and offers a scalable approach for next‑generation sequence modeling.

Abstract

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.

Paper Structure

This paper contains 33 sections, 11 equations, 4 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Preliminary
  4. Basic Notions and Consistency of Attention and SSM
  5. Efficiency of Attention and SSM with different token lengths
  6. Overall Framework of TransMamba
  7. Lossless Memory Converter
  8. Flexible TransPoint Scheduling
  9. Principles of TransPoint Scheduling
  10. Detailed TransPoint Schedule Designing
  11. Diverse Inference Strategy
  12. Experiments
  13. Experiment Setup
  14. Main Results
  15. Evaluations on General Tasks
  16. ...and 18 more sections

Figures (4)

  • Figure 1: TransMamba has shared parameters to flexibly switch between Attention and SSM, and TransPoints decide which parts of token sequence use Attention or SSM.
  • Figure 2: TransMamba generally shows better efficiency and performance with different sizes.
  • Figure 3: (a) Structure of TransMamba. Attention and SSM have shared parameters $\mathbf{W_{QKV}}$ and $\mathbf{W_{CBx}}$. Tokens are either processed via the green path (SSM mode) or the blue path (Attention mode). (b) Memory Converter. (c) The TransPoint Scheduling of TransMamba.
  • Figure (a): Layer-Shared TransPoint Schedule