¿Puede Mamba aprender a aprender? Un estudio comparativo sobre tareas de aprendizaje en 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

Resumen

Los modelos de espacio de estados (SSMs, por sus siglas en inglés), como Mamba Gu & Dao (2034), se han propuesto como alternativas a las redes Transformer en el modelado del lenguaje, al incorporar mecanismos de compuerta, convoluciones y selección de tokens dependiente de la entrada para mitigar el costo cuadrático de la atención multi-cabeza. Aunque los SSMs muestran un rendimiento competitivo, sus capacidades de aprendizaje en contexto (ICL, por sus siglas en inglés), una propiedad emergente notable de los modelos de lenguaje modernos que permite la ejecución de tareas sin optimización de parámetros, siguen siendo menos exploradas en comparación con los Transformers. En este estudio, evaluamos el rendimiento en ICL de los SSMs, centrándonos en Mamba, frente a los modelos Transformer en diversas tareas. Nuestros resultados muestran que los SSMs tienen un desempeño comparable al de los Transformers en tareas estándar de regresión en ICL, mientras que los superan en tareas como el aprendizaje de paridad dispersa. Sin embargo, los SSMs se quedan cortos en tareas que involucran funcionalidades de recuperación no estándar. Para abordar estas limitaciones, introducimos un modelo híbrido, \variant, que combina Mamba con bloques de atención, superando a los modelos individuales en tareas donde estos tienen dificultades por separado. Nuestros hallazgos sugieren que las arquitecturas híbridas ofrecen vías prometedoras para mejorar el ICL en los modelos de lenguaje.

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.