A Mamba Pode Aprender a Aprender? Um Estudo Comparativo sobre Tarefas de Aprendizado em Contexto

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

February 6, 2024

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

cs.AI

Resumo

Modelos de espaço de estados (SSMs), como o Mamba proposto por Gu & Dao (2034), foram sugeridos como alternativas às redes Transformer na modelagem de linguagem, incorporando mecanismos de gating, convoluções e seleção de tokens dependente da entrada para mitigar o custo quadrático da atenção multi-head. Embora os SSMs apresentem desempenho competitivo, suas capacidades de aprendizado em contexto (ICL), uma propriedade emergente notável dos modelos de linguagem modernos que permite a execução de tarefas sem otimização de parâmetros, permanecem pouco exploradas em comparação com os Transformers. Neste estudo, avaliamos o desempenho de ICL dos SSMs, com foco no Mamba, em comparação com modelos Transformer em diversas tarefas. Nossos resultados mostram que os SSMs têm desempenho comparável aos Transformers em tarefas padrão de regressão ICL, enquanto superam os Transformers em tarefas como aprendizado de paridade esparsa. No entanto, os SSMs apresentam deficiências em tarefas que envolvem funcionalidades de recuperação não padrão. Para abordar essas limitações, introduzimos um modelo híbrido, \variant, que combina o Mamba com blocos de atenção, superando os modelos individuais em tarefas onde eles têm dificuldades atuando de forma independente. Nossas descobertas sugerem que arquiteturas híbridas oferecem caminhos promissores para aprimorar o ICL em modelos de linguagem.

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.