Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning
  Tasks
Authors: name: Jongho Park; name: Jaeseung Park; name: Zheyang Xiong; name: Nayoung Lee; name: Jaewoong Cho; name: Samet Oymak; name: Kangwook Lee; name: Dimitris Papailiopoulos
Published: 2024-02-06 18:56:35

Abstract: 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.