Exploring the Limitations of Mamba in COPY and CoT Reasoning

TL;DR

This paper critically examines Mamba's capabilities in sequence modeling, revealing that while it can match Transformers in some tasks, it faces limitations in COPY and Chain of Thought reasoning, especially with fixed size.

Contribution

The paper provides a detailed analysis of Mamba's expressive power, highlighting its limitations and conditions under which it can or cannot outperform Transformers.

Findings

Mamba struggles with COPY operations at constant size.

Linear growth in Mamba size enables COPY but loses efficiency.

Mamba's performance on CoT tasks is limited compared to Transformers.

Abstract

Transformers have become the backbone of modern Large Language Models (LLMs); however, their inference overhead grows linearly with the sequence length, posing challenges for modeling long sequences. In light of this, Mamba has attracted attention for maintaining a constant inference size, with empirical evidence demonstrating that it can match Transformer performance in sequence modeling while significantly reducing computational costs. However, an open question remains: can Mamba always bring savings while achieving performance comparable to Transformers? In this paper, we focus on analyzing the expressive ability of Mamba to perform our defined COPY operation and Chain of Thought (CoT) reasoning. First, inspired by the connection between Mamba and linear attention, we show that constant-sized Mamba may struggle to perform COPY operations while Transformers can handle them more easily. However, when the size of Mamba grows linearly with the input sequence length, it can accurately perform COPY, but in this case, Mamba no longer provides overhead savings. Based on this observation, we further analyze Mamba's ability to tackle CoT tasks, which can be described by the Dynamic Programming (DP) problems. Our findings suggest that to solve arbitrary DP problems, the total cost of Mamba is still comparable to standard Transformers. However, similar to efficient Transformers, when facing DP problems with favorable properties such as locality, Mamba can provide savings in overhead. Our experiments on the copy and CoT tasks further demonstrate Mamba's limitations compared to Transformers in learning these tasks.