Le Mamba peut-il apprendre à apprendre ? Une étude comparative sur les tâches d'apprentissage en contexte

Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

February 6, 2024

Auteurs: Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, Dimitris Papailiopoulos

cs.AI

Résumé

Les modèles d'espace d'états (SSMs), tels que Mamba Gu & Dao (2034), ont été proposés comme alternatives aux réseaux Transformer dans la modélisation du langage, en intégrant des mécanismes de gating, des convolutions et une sélection de tokens dépendante de l'entrée pour atténuer le coût quadratique de l'attention multi-têtes. Bien que les SSMs affichent des performances compétitives, leurs capacités d'apprentissage en contexte (ICL), une propriété émergente remarquable des modèles de langage modernes permettant l'exécution de tâches sans optimisation des paramètres, restent moins explorées par rapport aux Transformers. Dans cette étude, nous évaluons les performances ICL des SSMs, en nous concentrant sur Mamba, par rapport aux modèles Transformer sur diverses tâches. Nos résultats montrent que les SSMs se comportent de manière comparable aux Transformers dans les tâches ICL de régression standard, tout en les surpassant dans des tâches comme l'apprentissage de parité sparse. Cependant, les SSMs sont moins performants dans les tâches impliquant une fonctionnalité de récupération non standard. Pour pallier ces limitations, nous introduisons un modèle hybride, \variant, qui combine Mamba avec des blocs d'attention, surpassant les modèles individuels dans les tâches où ils peinent indépendamment. Nos résultats suggèrent que les architectures hybrides offrent des perspectives prometteuses pour améliorer l'ICL dans les modèles de langage.

English

State-space models (SSMs), such as Mamba Gu & Dao (2034), have been proposed as alternatives to Transformer networks in language modeling, by incorporating gating, convolutions, and input-dependent token selection to mitigate the quadratic cost of multi-head attention. Although SSMs exhibit competitive performance, their in-context learning (ICL) capabilities, a remarkable emergent property of modern language models that enables task execution without parameter optimization, remain underexplored compared to Transformers. In this study, we evaluate the ICL performance of SSMs, focusing on Mamba, against Transformer models across various tasks. Our results show that SSMs perform comparably to Transformers in standard regression ICL tasks, while outperforming them in tasks like sparse parity learning. However, SSMs fall short in tasks involving non-standard retrieval functionality. To address these limitations, we introduce a hybrid model, \variant, that combines Mamba with attention blocks, surpassing individual models in tasks where they struggle independently. Our findings suggest that hybrid architectures offer promising avenues for enhancing ICL in language models.