TransMamba: Cambio flexible entre Transformer y 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

Resumen

Los Transformers son la piedra angular de los modelos de lenguaje modernos a gran escala, pero su complejidad computacional cuadrática limita la eficiencia en el procesamiento de secuencias largas. Los avances recientes en Mamba, un modelo de espacio de estados (SSM, por sus siglas en inglés) con complejidad lineal, ofrecen ganancias prometedoras en eficiencia, pero adolecen de inestabilidad en el aprendizaje contextual y la generalización multitarea. Este artículo propone TransMamba, un marco novedoso que unifica Transformer y Mamba a través de matrices de parámetros compartidas (por ejemplo, QKV y CBx), permitiendo así cambiar dinámicamente entre mecanismos de atención y SSM en diferentes longitudes de tokens y capas. Diseñamos el Conversor de Memoria para conectar Transformer y Mamba mediante la conversión de salidas de atención en estados compatibles con SSM, asegurando un flujo de información sin interrupciones en los TransPoints donde ocurre la transformación. También se explora a fondo la programación de TransPoints para obtener mejoras adicionales. Realizamos experimentos exhaustivos que demuestran que TransMamba logra una eficiencia de entrenamiento y un rendimiento superiores en comparación con los modelos base, y validamos la consistencia más profunda entre los paradigmas de Transformer y Mamba, ofreciendo una solución escalable para el modelado de secuencias de próxima generación.

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.