Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic data

TL;DR

The paper investigates the Achilles' heel of the Mamba architecture by using carefully designed synthetic tasks to reveal symmetry-related limitations. It shows that the nonlinear convolution before the State Space Model introduces an intrinsic asymmetry that biases Mamba toward composite solutions and impedes symmetric or reversed-sequence tasks, such as inverse sequence matching. Importantly, the root cause is the convolution stage rather than the SSM itself, and the authors demonstrate that architectural tweaks (e.g., residual connections and positional encoding) can mitigate these effects, aligning Mamba more closely with Transformer-like capabilities. These findings offer concrete guidance for designing future long-sequence models that retain linear complexity while improving symmetry-aware pattern recognition.

Abstract

State Space Models (SSMs) have emerged as promising alternatives to attention mechanisms, with the Mamba architecture demonstrating impressive performance and linear complexity for processing long sequences. However, the fundamental differences between Mamba and Transformer architectures remain incompletely understood. In this work, we use carefully designed synthetic tasks to reveal Mamba's inherent limitations. Through experiments, we identify that Mamba's nonlinear convolution introduces an asymmetry bias that significantly impairs its ability to recognize symmetrical patterns and relationships. Using composite function and inverse sequence matching tasks, we demonstrate that Mamba strongly favors compositional solutions over symmetrical ones and struggles with tasks requiring the matching of reversed sequences. We show these limitations stem not from the SSM module itself but from the nonlinear convolution preceding it, which fuses token information asymmetrically. These insights provide a new understanding of Mamba's constraints and suggest concrete architectural improvements for future sequence models.

Figures (17)

• Figure 1: Overview of Mamba and Transformer Blocks. The green trapezoids represent linear mappings. "smax" denotes the softmax function, "FNN" stands for feed-forward neural network, and "LN" represents layer normalization. The meanings of variables specific to the Mamba block are explained in the main text.
• Figure 2: Illustration of the Composite Function Task. Anchors 1, 2, 3, and 4 (depicted in orange) each represent distinct functions. Among the 16 anchor pairs formed, 14 correspond to composite functions derived directly from the sequential application of the individual anchor functions. The pair "34", highlighted in red, is defined as a different function rather than a direct composition. The pair "43" is intentionally excluded from the training set. The input to each anchor pair function is referred to as the "key" (indicated in green). Label indicates the output of an anchor pair applied to a key.
• Figure 3: Illustration of the Inverse Sequence Matching Task. The orange elements denote the generation set, which consists of three distinct numbers randomly selected from the interval [20, 100], as well as all possible permutations thereof. Blue and green indicate selected key sequences from the permutation space, separated by random numbers that do not belong to the generation set. One green key sequence is chosen as the answer sequence. The query sequence (shown in red) is obtained by reversing the answer sequence. The corresponding label identifies the position of the answer sequence. To prevent Mamba from leveraging its nonlinear convolution mechanism to infer answers, we prepend the query sequence with random numbers (outside the generation set) matching the length of Mamba's pure convolutional receptive field.
• Figure 4: Phase diagram of Mamba on the composite function task. Accuracy (color) for composite function task under different initialization rates (abscissa) and depths (ordinate). The groundtruth for (a) is composite solution and for (b) is symmetric solution. Detailed model configurations and training settings are provided in the appendix.
• Figure 5: SSM information flow in Mamba for the composite function task. Left: SSM information flow; green crosses indicate pruned connections. Right: (i) SSM input computation; (ii) state replacement after convolution. Flow is computed from $S = Mask\circ C^TB$, with line thickness indicating flow magnitude. Attention score from token $j$ to token $i$ is given by the $(i,j)$ entry of $S$ The numbers are the outputs of each layer through the model’s final linear layer and then take the arg‑max of the resulting logits to obtain the corresponding digit.