[논문 리뷰] Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

Jong-Ho Park, Jaeseung Park외 6명|arXiv (Cornell University)|2024. 02. 06.

Poverty, Education, and Child Welfare인용 수 10

한 줄 요약

이 논문은 상태 공간 모델(SSMs)에서의 인-context 학습(ICL)을 특히 Mamba를 Transformers와 비교하여 다양한 ICL 태스크에서 평가하고, 주의 기계 없는 ICL 능력과 주의 기반 능력을 결합한 하이브리드 MambaFormer를 도입하여 태스크 전반에서 강한 ICL 성능을 달성한다.

ABSTRACT

State-space models (SSMs), such as Mamba (Gu & Dao, 2023), 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, MambaFormer, 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.

연구 동기 및 목표

• 다양한 태스크 전반에서 SSMs, 특히 Mamba가 처음부터 인-context 학습(ICL)을 학습할 수 있는지 평가한다.
• 표준 ICL 태스크에서 Mamba를 Transformer 모델과 비교한다. 태스크에는 regression, parity, retrieval, 및 chain-of-thought가 포함된다.
• 하이브리드 아키텍처가 SSM과 주의 메커니즘의 강점을 활용해 ICL을 향상시킬 수 있는지 조사한다.
• retrieval 및 의사결정 트리( decision-tree) 태스크에서 Mamba의 한계점을 탐구하고 이러한 격차를 완화하는 아키텍처를 확인한다.

제안 방법

• 회귀, 이상치, 이산 함수, chain-of-thought, 및 retrieval를 포괄하는 ICL 태스크 모음에서 처음부터 학습한다.
• 또한 positional encoding이 있는/없는 Transformer와 두 하이브리드인 Standard Hybrid 및 MambaFormer를 포함해 Mamba, S4, S4-Mamba, Transformer를 평가한다.
• Mamba 블록을 주의 블록과 교대로 배치하여 하이브리드 아키텍처를 형성하고 태스크 간 성능을 비교한다.
• 고정된 학습 목표를 사용한다: 태스크마다 임의로 생성된 프롬프트와 입력에 대해 태스크 특이 손실을 최소화한다.
• 태스크당 1,280 개의 테스트 프롬프트에 걸친 경험적 손실과 수렴 동작 측면에서 성능을 보고한다.

Figure 1 : MambaFormer is a hybrid architecture that replaces MLP blocks within the transformer with Mamba blocks. Importantly, the architecture also starts with a Mamba block and does not use positional encoding. In our ICL evaluations, we find that MambaFormer consistently achieves a best-of-both-

실험 결과

연구 질문

• RQ1 본 연구에서 탐구된 규모에서 Mamba 및 다른 SSM이 in-context 작업을 학습할 수 있는가?
• RQ2 표준 ICL 태스크에서 Mamba는 Transformer 모델과 어떻게 비교되는가(linear, sparse, 2NN, 의사결정 트리, outliers, sparse parity, CoT I/O, MQAR)?
• RQ3 하이브리드 아키텍처(Standard Hybrid, MambaFormer)가 태스크 전반에서 양쪽의 장점을 살린 최상의 성능을 제공하는가?
• RQ4 Mamba 및 Transformer의 한계를 드러내는 태스크는 무엇이며, 하이브리드가 이를 완화할 수 있는가?

주요 결과

• Mamba는 처음부터 학습시켜 다양한 ICL 태스크를 수행하게 할 수 있으며 표준 회귀 태스크에서 Transformer 성능에 종종 근접한다.
• Mamba는 특히 복잡한 ICL 태스크(예: sparse parity, 긴 시퀀스의 outlier 회귀)에서 Transformer가 어려움을 겪는 경우에도 뛰어나지만 일부 retrieval 및 decision-tree 태스크에서는 실패한다.
• Transformers는 decision-tree 및 벡터-valued MQAR retrieval 태스크에서 Mamba보다 우수한 반면, Mamba는 특정 retrieval 및 parity 관련 태스크에서 Transformer를 능가한다.
• 하이브리드 아키텍처, 특히 MambaFormer (Mamba 블록 뒤에 attention이 오는 구성, positional encoding 없음)은 전체 태스크 모음에서 강력한 성능을 달성하며 두 모델의 강점을 결합한다.
• Mamba를 초기 층으로 두고 Mamba와 attention 레이어를 교대로 배치하는 것이 parity 학습 및 향상된 retrieval 능력을 가능하게 하는 데 중요하다.
• MambaFormer는 대부분의 태스크에서 Transformer 및 Mamba에 필적하거나 이를 능가할 수 있으며, 고잡음 환경에서도 효율적으로 선형 회귀를 학습하여 전반적인 ICL 능력을 보여준다.