Les Transformers sont des SSM : Modèles Généralisés et Algorithmes Efficaces via la Dualité des Espaces d'État Structurés

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

May 31, 2024

Auteurs: Tri Dao, Albert Gu

cs.AI

Résumé

Alors que les Transformers ont été l'architecture principale derrière le succès de l'apprentissage profond en modélisation du langage, les modèles à espace d'états (SSMs) tels que Mamba ont récemment démontré des performances équivalentes ou supérieures aux Transformers à petite et moyenne échelle. Nous montrons que ces familles de modèles sont en réalité étroitement liées, et développons un cadre théorique riche de connexions entre les SSMs et des variantes de l'attention, reliées à travers diverses décompositions d'une classe bien étudiée de matrices semi-séparables structurées. Notre cadre de dualité d'espace d'états (SSD) nous permet de concevoir une nouvelle architecture (Mamba-2) dont la couche centrale est un raffinement du SSM sélectif de Mamba, 2 à 8 fois plus rapide, tout en restant compétitive avec les Transformers en modélisation du langage.

English

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.