TransMamba: 트랜스포머와 맘바 간 유연한 전환

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

초록

트랜스포머는 현대의 대규모 언어 모델의 핵심이지만, 이차원적 계산 복잡도로 인해 긴 시퀀스 처리에서 효율성이 제한됩니다. 선형 복잡도를 가진 상태 공간 모델(SSM)인 Mamba의 최근 발전은 효율성 향상을 약속하지만, 불안정한 문맥 학습과 다중 작업 일반화 문제를 겪고 있습니다. 본 논문은 공유 매개변수 행렬(예: QKV 및 CBx)을 통해 트랜스포머와 Mamba를 통합한 새로운 프레임워크인 TransMamba를 제안합니다. 이를 통해 다양한 토큰 길이와 레이어에서 주의 메커니즘과 SSM 메커니즘 사이를 동적으로 전환할 수 있습니다. 우리는 주의 출력을 SSM 호환 상태로 변환하여 변환이 발생하는 TransPoints에서 원활한 정보 흐름을 보장하는 Memory converter를 설계했습니다. 또한 TransPoint 스케줄링을 철저히 탐구하여 추가 개선을 이루었습니다. 광범위한 실험을 통해 TransMamba가 기준 모델 대비 우수한 학습 효율성과 성능을 달성함을 입증했으며, 트랜스포머와 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.