TransMamba: Alternando de Forma Flexível entre Transformer e Mamba

TransMamba: Flexibly Switching between Transformer and Mamba

March 31, 2025

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

cs.AI

Resumo

Os Transformers são a base dos modelos de linguagem de grande escala modernos, mas sua complexidade computacional quadrática limita a eficiência no processamento de sequências longas. Avanços recentes no Mamba, um modelo de espaço de estados (SSM) com complexidade linear, oferecem ganhos promissores de eficiência, mas sofrem com aprendizado contextual instável e generalização multitarefa. Este artigo propõe o TransMamba, uma nova estrutura que unifica o Transformer e o Mamba por meio de matrizes de parâmetros compartilhadas (por exemplo, QKV e CBx), permitindo alternar dinamicamente entre mecanismos de atenção e SSM em diferentes comprimentos de tokens e camadas. Projetamos o Conversor de Memória para conectar o Transformer e o Mamba, convertendo saídas de atenção em estados compatíveis com SSM, garantindo um fluxo de informações contínuo nos TransPoints onde a transformação ocorre. O agendamento de TransPoints também é explorado detalhadamente para melhorias adicionais. Realizamos extensos experimentos que demonstram que o TransMamba alcança eficiência de treinamento e desempenho superiores em comparação com as abordagens de referência, validando a consistência mais profunda entre os paradigmas do Transformer e do Mamba, oferecendo uma solução escalável para a modelagem de sequências de próxima geração.

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.