Mamba può imparare ad imparare? Uno studio comparativo sui compiti di apprendimento in contesto

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

February 6, 2024

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

cs.AI

Abstract

I modelli spazio-stato (SSM), come Mamba Gu & Dao (2034), sono stati proposti come alternative alle reti Transformer nel campo del modeling linguistico, incorporando meccanismi di gating, convoluzioni e selezione di token dipendente dall'input per mitigare il costo quadratico dell'attenzione multi-testina. Sebbene gli SSM mostrino prestazioni competitive, le loro capacità di apprendimento in contesto (ICL), una proprietà emergente notevole dei moderni modelli linguistici che consente l'esecuzione di task senza ottimizzazione dei parametri, rimangono meno esplorate rispetto ai Transformer. In questo studio, valutiamo le prestazioni ICL degli SSM, con un focus su Mamba, confrontandoli con i modelli Transformer su vari task. I nostri risultati mostrano che gli SSM performano in modo comparabile ai Transformer nei task ICL di regressione standard, superandoli in task come l'apprendimento della parità sparsa. Tuttavia, gli SSM risultano carenti in task che coinvolgono funzionalità di recupero non standard. Per affrontare queste limitazioni, introduciamo un modello ibrido, \variant, che combina Mamba con blocchi di attenzione, superando i modelli individuali nei task in cui faticano separatamente. Le nostre scoperte suggeriscono che le architetture ibride offrono percorsi promettenti per migliorare l'ICL nei modelli linguistici.

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.