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

Аннотация

Трансформеры являются краеугольным камнем современных крупных языковых моделей, однако их квадратичная вычислительная сложность ограничивает эффективность при обработке длинных последовательностей. Недавние достижения в модели Mamba, основанной на пространстве состояний (SSM) с линейной сложностью, предлагают перспективные улучшения в эффективности, но страдают от нестабильного контекстного обучения и обобщения на множественные задачи. В данной статье предлагается TransMamba — новая архитектура, которая объединяет Transformer и Mamba через общие матрицы параметров (например, QKV и CBx), что позволяет динамически переключаться между механизмами внимания и SSM на разных уровнях и длинах токенов. Мы разработали Memory converter для соединения Transformer и Mamba путем преобразования выходов внимания в состояния, совместимые с SSM, обеспечивая плавный поток информации в точках трансформации (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.