Может ли Mamba научиться учиться? Сравнительное исследование задач обучения в контексте

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

February 6, 2024

Авторы: Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, Dimitris Papailiopoulos

cs.AI

Аннотация

Модели пространства состояний (SSMs), такие как Mamba Gu & Dao (2034), были предложены в качестве альтернативы сетям Transformer в задачах языкового моделирования. Они включают механизмы гейтирования, свертки и зависимый от входных данных выбор токенов, чтобы смягчить квадратичную сложность многоголового внимания. Хотя SSMs демонстрируют конкурентоспособную производительность, их способности к обучению в контексте (ICL) — замечательное эмерджентное свойство современных языковых моделей, позволяющее выполнять задачи без оптимизации параметров, — остаются недостаточно изученными по сравнению с Transformers. В данном исследовании мы оцениваем производительность SSMs, в частности Mamba, в задачах ICL, сравнивая их с моделями Transformer на различных задачах. Наши результаты показывают, что SSMs демонстрируют сопоставимую с Transformers производительность в стандартных задачах регрессии ICL, превосходя их в задачах, таких как обучение разреженной четности. Однако SSMs уступают в задачах, связанных с нестандартной функциональностью извлечения данных. Для устранения этих ограничений мы представляем гибридную модель \variant, которая объединяет Mamba с блоками внимания, превосходя отдельные модели в задачах, где они испытывают трудности по отдельности. Наши результаты указывают на то, что гибридные архитектуры открывают перспективные пути для улучшения ICL в языковых моделях.

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.