Kan Mamba Leren Hoe te Leren? Een Vergelijkende Studie over In-Context Leren Taken

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

Samenvatting

State-space modellen (SSM's), zoals Mamba Gu & Dao (2034), zijn voorgesteld als alternatieven voor Transformer-netwerken in taalmodellering, door het incorporeren van gating, convoluties en invoerafhankelijke tokenselectie om de kwadratische kosten van multi-head attention te verminderen. Hoewel SSM's competitieve prestaties vertonen, blijven hun in-context learning (ICL)-mogelijkheden, een opmerkelijke emergente eigenschap van moderne taalmodellen die taakuitvoering zonder parameteroptimalisatie mogelijk maakt, onderbelicht in vergelijking met Transformers. In deze studie evalueren we de ICL-prestaties van SSM's, met de focus op Mamba, tegenover Transformer-modellen in verschillende taken. Onze resultaten laten zien dat SSM's vergelijkbaar presteren met Transformers in standaard ICL-regressietaken, terwijl ze hen overtreffen in taken zoals sparse parity learning. SSM's presteren echter minder goed in taken die niet-standaard retrieval-functionaliteit vereisen. Om deze beperkingen aan te pakken, introduceren we een hybride model, \variant, dat Mamba combineert met attention-blokken, en dat individuele modellen overtreft in taken waarin ze afzonderlijk tekortschieten. Onze bevindingen suggereren dat hybride architecturen veelbelovende mogelijkheden bieden om ICL in taalmodellen te verbeteren.

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.