TransMamba: Passaggio flessibile tra Transformer e Mamba

TransMamba: Flexibly Switching between Transformer and Mamba

March 31, 2025

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

cs.AI

Abstract

I Transformer sono il fondamento dei moderni modelli linguistici su larga scala, ma la loro complessità computazionale quadratica limita l'efficienza nell'elaborazione di sequenze lunghe. I recenti progressi in Mamba, un modello a spazio di stati (SSM) con complessità lineare, offrono promettenti guadagni in termini di efficienza, ma soffrono di un apprendimento contestuale instabile e di una generalizzazione multitask. Questo articolo propone TransMamba, un nuovo framework che unisce Transformer e Mamba attraverso matrici di parametri condivise (ad esempio, QKV e CBx), consentendo così di passare dinamicamente tra meccanismi di attenzione e SSM a diverse lunghezze di token e livelli. Progettiamo il Memory converter per collegare Transformer e Mamba convertendo gli output di attenzione in stati compatibili con SSM, garantendo un flusso di informazioni senza soluzione di continuità nei TransPoint dove avviene la trasformazione. Viene inoltre esplorata a fondo la pianificazione dei TransPoint per ulteriori miglioramenti. Abbiamo condotto esperimenti estensivi dimostrando che TransMamba raggiunge un'efficienza di addestramento e prestazioni superiori rispetto ai baseline, e abbiamo validato una più profonda coerenza tra i paradigmi di Transformer e Mamba, offrendo una soluzione scalabile per la modellazione di sequenze di prossima generazione.

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.