Can Mamba Always Enjoy the "Free Lunch"?

Q2: Similarly, the analysis of the paper depends on critical Assumptions 1 and 2. While I understand these are mathematical constructions needed to make the derivations go through, these assumptions appear to be a bit too convoluted and difficult be measured or justified empirically. In general, there are many ways to derive theoretical upper and lower bounds, depending on the choice of assumptions, however an “optimal” theory should be based on minimal assumptions while producing interpretable outcomes.

First, regarding Assumption 1: for the first condition, the appearance of $\rho$ and $\delta$ serves primarily as a convenient definition for expressing the theorem. We do not impose strict restrictions on the values of $\rho$ and $\delta$. Their sizes are flexible, but they generally form a trade-off with the conditions in Theorem 1; the second condition is simply to ensure the stability of the output scale. This assumption is necessary; otherwise, during theoretical analysis, the output scale could increase significantly as $|S_i|$ grows. In practice, this issue can be mitigated by using modules like layer normalization. However, incorporating such mechanisms into theoretical analysis would introduce additional complications, which is why we opted for the current assumption.

We would like to provide an example: Given a text snippet, "XXXXXXX(Some unrelated content). Bob has a dog while John has a cat. After school, Bob often plays with his ____." We assume the expected output from the model here is "dog". In this example, assuming its position of the answer is $i$, the Matching Set $S_i$ for the current token to be predicted and $L \approx 18$ represents the distance within which the answer might exist relative to position $i$. In other words, starting from the phrase "Bob has a ...". In addition, $pos(i)$ refers to the position index of the word "dog". The Matching Set $S_i$ might include position indices for words like "Bob," "dog," "John," and "cat" if we choose some $\delta$. Of course, if we choose a larger $\delta$, the Matching Set $S_i$ might only include the position indices corresponding to "Bob" and "dog".

For Assumption 2: A significant number of DP problems conform to such conditions. For example, consider the problem of finding the minimum edit distance, where we aim to transform string $s_1$ into string $s_2$ with lengths $n_1 = |s_1|$ and $n_2 = |s_2|$, respectively. The costs for insertion, deletion, and substitution are $a$, $b$, and $c$, respectively. Let $\rm dp(j,k)$ represent the cost of transforming the first $j$ characters of $s_1$ into the first $k$ characters of $s_2$. The transition function can then be expressed as:

Finally, the aggregation function selects $\rm dp(n_1, n_2)$ as the final answer. In the example above, the size of the state space, all intermediate values, and the lengths of the input strings will all be upper-bounded by $\rm poly(n_1, n_2)$. Moreover, the operations required by the involved functions can be approximated with polynomial efficiency by a constant-size MLP, as shown in Lemmas 1-5. Therefore, such DP problems satisfy Assumption 2.

Q3: It would be great to also provide some experimental examples to support the claim of Sec. 5 (Mamba does not offer additional cost savings compared to Transformers for DP problems), similar to the COPY problem.

Thank you for your suggestion. We will design CoT experiments in the future revisions to support our claim.

Q4: In Fig. 2, what is the Transformer baseline in the rightmost chart?

In Fig. 2, as mentioned in Section 4.3.2, we did not compare against Transformer; instead, we only compared different sizes of pre-trained Mamba blocks.

We greatly appreciate the reviewer's careful review and constructive suggestions. Below are our responses.

W1: Existing works already demonstrated similar limitations of Mamba (and SSM in general) in handling tasks like COPY (especially [1,2], which the paper also referred to). The conclusions of this paper did not appear to have provided much additional insight, and the contributions are incremental.

The tasks we focus on are fundamentally different: Merrill et al. (2024) [1] emphasize that Mamba is also unable to solve state tracking; Jalessi et al. (2024) [2] concentrate on scaling in terms of sequence length and the number of states. In contrast, our work examines the impact of the distance to the target token on output error from the perspective of numerical approximation. Additionally, our underlying intuition is distinct: our analysis is inspired by the relationship between linear attention and Mamba, which is an aspect not addressed in the two aforementioned works.

[1] The Illusion of State in State-Space Models

[2] Repeat After Me: Transformers are Better than State Space Models at Copying

W2: What is more, the theoretical derivation and results of the paper (Theorem 1-4, and the assumptions they build upon) seem to be pure complex mathematical constructions, which do not appear to be concise and interpretable enough to be fully convincing or empirically verifiable.

Intuitively, in Theorem 1, we provided a sufficient condition for Mamba to perform the COPY operation, while Theorems 2 to 4 aim to establish an upper bound on the model size required by Mamba blocks to solve the target tasks. That is, they show the maximum model size needed to accomplish COPY/DP tasks.

In addition, as highlighted by reviewer k7wk, we will include additional examples in future versions to make the implications of these conclusions more intuitive (you can also see our response to Q2). Furthermore, similar to the theoretical analysis in prior works, our theoretical section focuses on the expressive power of the Mamba block, rather than whether this ability can be learned. However, in Section 4.3, we provide some experiments to demonstrate phenomena related to these theoretical results.

Q1: The theoretical result in Theorem 1 is honestly a bit difficult to understand intuitively: Is this a loose or tight bound? It appears to be a characterization of the difficulty of a COPY problem, but it is difficult to get a sense of what it indicates in practice as there are many variables introduced to control this bound. It would be great to justify the interpretability of this theoretical result and how this bound is an accurate characterization of the model behavior (e.g. plot curves comparing theoretical bounds with empirical results).

This explanation of Theorem 1 might make it easier to understand: for constant-size Mamba blocks, we identified the sufficient conditions under which Mamba can perform the COPY operation. Rather than emphasizing whether these conditions form a tight or loose bound, we prefer to describe them as stringent. This is stringent because we are considering a fairly general scenario involving a given sequence $X = [x_i]_{i=1}^{N}$ and positions $pos(i)$. Of course, Mamba blocks may still be able to perform the COPY operation in less restrictive cases where the sufficient conditions are not satisfied. That is, given any sequence and position, it remains unclear whether constant-size Mamba blocks can always perform the COPY operation. However, we want to point out that Mamba blocks of size $O(N)$ always can.

I'd like to thank authors for your responses. I would like to keep my original rating.

We greatly appreciate the reviewer's careful review and constructive suggestions. Below are our responses.

W1: While the inclusion of DP problems feels necessary to relate some of the abstract findings to a real problem (CoT with DP), the findings feel inconclusive because it teaches us about representational power and is an existence proof. However, as we observe in the Phonebook experiments, that doesn’t translate in practice to viability; in particular the bounds are still quite loose. When reading the paper, it felt like Section 5 was a separate paper as it introduces a new set of notation despite the result being ultimately a corollary of the original theorem.

First, regarding these "abstract findings" about "representational power", we would interpret them as conclusions about the upper bound on the computational overhead required for the model to complete the task. Specifically, through constructive proofs, we provide an upper bound on the computational cost Mamba incurs when using CoT to solve DP tasks. Furthermore, concerning the distinction between "abstract findings" and "real problems," we believe adding more illustrative examples could help enhance the reader's understanding, as referenced in our response to W2. In fact, our initial intention was to follow the techniques of Feng et al. [1] to investigate the improvements in representational power brought by Mamba under CoT (Theorem 3 in ours and Theorem 4.7 in Feng et al. [1]). To explore Mamba block's ability to address DP problems, we required several lemmas (Theorem 2, Lemma 1-5) to demonstrate the fundamental operations Mamba blocks can perform (e.g., multiplication, approximate MLPs, linear transformations). During this process, we discovered that the COPY operation is critical for solving DP problems. Moreover, Mamba and Transformers exhibit significant differences in executing this operation, which led us to decide to dedicate a separate section to discussing Mamba's ability to perform the COPY operation.

Therefore, the content in Section 4 is closely related to Section 5. Perhaps we should emphasize the differences between the two models in performing the COPY operation in Section 4, which could improve the overall organization of the content, as mentioned in our response to Q1.

[1] Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

Q1: “free lunch” is not well-defined, ... Do you agree with this generalized statement? ... not misunderstanding the paper.

What we intended to express with "Mamba always enjoys the free lunch" is that Mamba can accomplish a given task with a smaller model size (and inference time) compared to Transformers. To challenge this claim, in Section 4, we specifically focused on the defined COPY task, and showed the potential limitations of constant-sized models in reproducing historical information. In Theorem 1, we presented a sufficient condition for performing the COPY operation and pointed out that this condition is relatively difficult to satisfy. In such cases, Mamba cannot enjoy the free lunch.

We believe the ambiguity in the "free lunch" statement might also stem from the lack of a direct comparison with Transformers in this section. In fact, whether Transformers can perform the COPY task is partly explained by Lemma C.7 of Feng et al., 2023 [1] which demonstrates that a constant-sized attention layer (relative to sequence length) can achieve COPY-like behavior, satisfying $\| o_i - v_{pos(i)} \| \le \epsilon$. However, for any given token, this operation incurs linear time overhead. Similarly, for Mamba to achieve the COPY task, it might also require a linear model size and corresponding linear inference time. From this perspective, Mamba does not enjoy the free lunch for this task. However, this comparison remains somewhat rough, as the formalized assumptions differ between the two models (e.g., Assumption 1 in our work for Mamba and Assumption C.6 in Feng et al. for attention layers), though the underlying logic bears similarities.

In Section 5, we shifted focus to Mamba equipped with CoT, analyzing its upper bounds on time and space overhead when solving DP tasks compared to Transformers. This serves to show that Mamba cannot always enjoy the free lunch.

We agree with your comment, "If the distance ... recall that information," and believe your understanding is accurate.

Q2: What’s the significance of the $l_{\infty}$ norm in Theorem 2? Is it a corollary, assumption, or not actually relevant to the results?

It can be understood as a corollary, where all parameters being upper bounded by $\rm O(poly(M, N))$ implies that the COPY task can also be solved by the Mamba block with $\rm log(N)$ precision (Line 310). For the reason behind using a $\rm log$-precision model, please refer to Appendix A.3 of Feng et al., 2023 [1]. Briefly, a model with $\rm log(N)$ precision can represent real numbers bounded by $\rm O(poly(n))$ with a truncation error of $\rm O(poly(1/n))$.

Q3: Is it possible to express the DP conditions of Assumption 2 in terms of a formal problem complexity class...including a (admittedly large) subset of commonly-seen DP problems.

Assumption 2 follows the assumptions established by Feng et al., 2023 [1] and Yang et al., 2024 [2]. Under these same assumptions for DP problems, we can compare Mamba and Transformers. In the work of Merrill et al., 2023 [3], the capability upper bound of log-precision Transformers is $\rm TC^0$. Feng et al. demonstrated that Transformer+CoT can surpass this upper bound, but it still cannot solve the CFG problem (Theorem 4.8), which is $\rm P$-complete. Therefore, we intuitively believe that the problems described by Assumption 2 may lie between $\rm TC^0$ and $\rm P$-complete.

[1] Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

[2] Do Efficient Transformers Really Save Computation?

[3] The Parallelism Tradeoff: Limitations of Log-Precision Transformers

Q4: For Theorem 2: is it possible to give a tighter bound than poly(M,N)? L817 suggests it is NM2, is that a correct bound?

As mentioned in our response to Q2, the parameter upper bound being $\rm poly(M, N)$ is intended to show that a $\rm log(N)$-precision Mamba block can perform this operation. While we are unsure if there exists a tighter bound, a $\rm log(N)$-precision Mamba will suffice in any case. We are unclear about your question regarding "is that a correct bound." Are you suggesting that parameter sizes can be bounded by $\rm NM^2$, and that’s why we described it as $\rm poly(M, N)$? If so, that is exactly what we intended to convey.

C5: Sec 4.1, when we set H_0=O, what is O? Do you mean 0?

Yes, $\mathbf{O}$ means the matrix whose elements are all zeros and shape is same as $H_0$.

C6: L212 Why is a_i \in [0,1]? I know why it is positive, but it isn’t clear to me why it’s <=1

Just as inllustrated in Line 208 & 209, all elements of $A$ are usually set to be negative while $\Delta>0$, thus $\mathrm{exp}(\Delta A) \le 1$.

C1-C4 & C7-C13: Advices and typos

Thank you for your constructive suggestions and for carefully pointing out the typos. We will carefully adopt your recommendations and address these issues in subsequent versions.

W2： It isn’t stated clearly why the assumptions (both Assumption 1 and Assumption 2) are safe to assume. When are these assumptions satisfied?

Assumption 1, first condition – I don’t understand why it is reasonable to assume this...

Assumption 1, second condition - this is missing a something. “There exists Si such that…” or “There exists j such that?

For Assumption 1: First, in the first condition, the introduction of $\rho$ and $\delta$ serves more as a convenient definition for the theorem's formulation (perhaps it should not have been presented as an assumption). We do not impose restrictions on the values of $\rho$ and $\delta$​, allowing them to be relatively flexible.

For the second condition, as mentioned in Line 269, it is to ensure the stability of the output scale. This assumption is necessary; otherwise, during theoretical analysis, the output scale could grow significantly with the increase of $|S_i|$. In practice, this issue can be mitigated by using modules like layer normalization. However, such approaches would introduce additional complexity into the theoretical analysis. Hence, we adopted the current assumption.

We believe there might be some misunderstanding regarding the issues mentioned in Assumption 1, which has led to us not fully grasping the point in the question. To clarify this, we would like to provide an example: Given a text snippet, "XXXXXXX(Some unrelated content). Bob has a dog while John has a cat. After school, Bob often plays with his ____." We assume the expected output from the model here is "dog".

In this example, assuming its position of the answer is $i$, the Matching Set $S_i$ for the current token to be predicted and $L \approx 18$ represents the distance within which the answer might exist relative to position $i$. In other words, starting from the phrase "Bob has a ...,". In addition, $pos(i)$ refers to the position index of the word "dog". The Matching Set $S_i$ might include position indices for words like "Bob," "dog," "John," and "cat" if we choose some $\delta$. Of course, if we choose a larger $\delta$, the Matching Set $S_i$ might only include the position indices corresponding to "Bob" and "dog". Indeed, the suggestion that $\rho \ge \delta + \epsilon$ is a good idea for ensuring the "gap," and we will consider reorganizing and rephrasing this part in the future.

For Assumption 2: A significant number of DP problems conform to such conditions. For example, consider the problem of finding the minimum edit distance, where we aim to transform string $s_1$ into string $s_2$ with lengths $n_1 = |s_1|$ and $n_2 = |s_2|$, respectively. The costs for insertion, deletion, and substitution are $a$, $b$, and $c$, respectively. Let $\rm dp(j,k)$ represent the cost of transforming the first $j$ characters of $s_1$ into the first $k$ characters of $s_2$. The transition function can then be expressed as:

Finally, the aggregation function selects $\rm dp(n_1, n_2)$ as the final answer. In the example above, the size of the state space, all intermediate values, and the lengths of the input strings will all be upper-bounded by $\rm poly(n_1, n_2)$. Moreover, the operations required by the involved functions can be approximated with polynomial efficiency by a constant-size MLP, as shown in Lemmas 1-5. Therefore, such DP problems satisfy Assumption 2.

W3: The presentation of the core theorems (Section 4.2) is very difficult to understand. It might be easier to present Theorem 1 in terms of δ (and use ρ only in the proof - the first condition doesn't actually need ρ either). For Theorem 2, the “simple intuition” could be better expressed in relation to the terms of Eq. 7, i.e. as a direct result of the first term increasing.

Thank you for your suggestion. We have also noticed that the $\rho$ in Assumption 1 and Theorem 1 is redundant, and directly using $c_i^T b_{pos(i)}$ would indeed be more straightforward. Regarding "as a direct result of the first term increasing," we do not understand your meaning as Theorem 1 deals with a constant-size setting relative to $N$ and the first term does not seem to reflect any explicit relationship with $N$.

W4: Finally, the empirical results would be significantly more interpretable if instead of the total number of parameters, the actual values of the Mamba block size were reported, since that's what the theoretical results are in reference to.

We agree with your suggestion. Adding charts to illustrate the relationship between different block sizes (for example, 10M, 20M, ..., 130M, ...) and the number of training examples needed (just like Line 359) would better align with our theoretical results.

W3: Finally, this draft requires improvement in writing quality. The notation system is confusing; for example, it is unclear why in line 237, $\alpha_j= \Pi_{k=j+1}^i a_k$ is defined without any subscript or superscript in $\alpha$.

Thank you for your suggestions. We will further refine our expression in the future to avoid causing confusion or misunderstandings for the readers. For the example you provided, $\alpha$ is simply a constant introduced for simplicity (using $\Pi_{k=j+1}^i a_k$ would seem redundant), and it carries no additional meaning.

Q: Regarding the COPY task setting, it appears that for each token i, a "relevant token" $\rm pos(i)$ is randomly selected from a set $S_i$, defined such that the absolute value of $c_i^T b_{pos(i)}$ exceeds a threshold. Could you clarify the rationale behind this definition? I find it unclear why the COPY operation is structured this way

This definition is based on the attention perspective, meaning that when we want to recover historical records at some position $\rm pos(i)$ based on $c_i$ (query), the query ($c_i$) should maintain a sufficiently large attention score $c_i^T b_{pos(i)}$ with the key ($b_{pos(i)}$) corresponding to that position. This is also the difference in the starting point compared to the aforementioned work.

We are grateful for the time and effort the reviewer have dedicated to our paper. Below are our responses.

W1: Firstly, the theoretical contributions appear trivial... Stronger theoretical results on the state size requirements for RNNs (including Mamba) for recall tasks...as the proofs align closely with this existing body of work.

We believe there are significant differences here. Arora et al. [1] focus on a synthetic AR task called Multi-Query Associative Recall (MQAR), using communication complexity theory to show that recurrent models require at least $\Omega(N)$ to solve MQAR (Theorem 3.1). This provides a lower bound guarantee for solving the problem. In contrast, our Theorem 1 presents a sufficient condition for solving the COPY operation, and Theorem 2 provides an upper bound on the model size required to solve the COPY operation, meaning we use a constructive approach to show that the model size needed to complete the COPY task is at most $O(N)$.

Similarly, in Wen et al.'s Theorem 4.6 [2], they express that even with CoT, any RNN model with $o(n)$ bit memory cannot solve tasks in $T \in \{\text{Index, AR, c-gram retrieval, Counting}\}$ of size $n$ for large enough $n$. They also use communication complexity and information-theoretic arguments to prove the lower bound. In comparison, our Theorem 2 does not involve the use of CoT for the Mamba block. Therefore, we believe the two conclusions are not in conflict.

Additionally, Jalessi et al.[3] focus on scaling in terms of sequence length and the number of states, whereas our work focuses on analyzing the impact of the distance of the token to be copied on the output error from the numerical approximation perspective.

Finally, regarding the statement "the proofs align closely with this existing body of work", we believe there has been a significant misunderstanding. As already explained, the works mentioned above mainly rely on communication complexity theory or information-theoretic results, while our approach is grounded in numerical approximation. We have not found direct connection between our proofs and the existing proofs in these works. Therefore, we believe our theoretical contributions are distinct and do not closely align with the existing bodies of works.

[1] Simple linear attention language models balance the recall-throughput tradeoff

[2] RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval

[3] Repeat After Me: Transformers are Better than State Space Models at Copying

W2: Secondly, regarding the CoT proof, Mamba can be formulated as a (gated) linear attention model, as detailed ... the paper’s contribution remains ambiguous.

Undoubtedly, there is a growing body of work linking Mamba with linear attention, such as the work by Yang et al. [1] mentioned, and those we referenced in Section 2 of our paper. In Section 3, we restate Han et al.'s work [2] simply to provide clarity and context for our readers, as also agreed upon by reviewer PWf8. This section is, of course, not our original contribution, as we have emphasized multiple times in the text, and it was not also presented as contribution in the introduction.

Regarding the claim that we are "applying known results on Efficient Transformers (i.e., standard linear attention) from Wen et al. (2024) without adding unique insights," we believe this is a misunderstanding. In fact, in the CoT section, we follow the setup from Feng et al. [3] focusing on Mamba's ability to solve the DP problem and comparing the results with those of standard Transformer [3] and efficient Transformers [4]. Our conclusions in this part are based on Theorem 2 and we did not rely on the conclusions from Wen et al. [5]. Our Theorem 3 provides an upper bound on the model size required for Mamba to solve the DP problem when equipped with CoT. This is distinct from Wen et al.'s Theorem 4.6, which provides a lower bound and is concerned with different problem settings. Therefore, we believe there has been a misunderstanding about our contributions and the used approach.

[1] Gated Linear Attention Transformers with Hardware-Efficient Training

[2] Demystify mamba in vision: A linear attention perspective.

[3] Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

[4] Do Efficient Transformers Really Save Computation?

[5] RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval

W3: To me, the findings are not novel and not distinguished from existing works. Most of the paper focuses on introducing existing works.

We have also noticed that due to insufficient explanation on our part, the conclusions of our work "seem somewhat similar" to those of other works, and the distinctions are not clearly highlighted. In fact, our work differs from existing works in terms of the tasks we focus on, the underlying intuition, and the theoretical results. Reviewer vUfm in Weakness 1, BtyZ in Weakness 1&2, and Reviewer k7wk in Minor Comment 1 also expressed similar concerns about the novelty of our work. We suspect that the "existing works" include the ones mentioned in their comments, and this part can refer to our corresponding responses. We will add a more detailed comparison with these works in the next version.

W4: The main claim about the bottleneck of MAMBA is not discussed in this paper. While we can see this from empirical studies, it is not clear how we can take inspiration from the assumption the authors make. This is important as the authors claim they explain the limitations from theoretical standpoints.

As stated in our response to W2, the bottleneck logic we aim to express for Mamba is as follows: for Mamba blocks of constant size, we have identified sufficient condition for executing the COPY operation. We want to emphasize that satisfying the sufficient condition is stringent. Of course, even in more relaxed situations where the sufficient conditions are not met, Mamba blocks may still be able to perform the COPY operation. The question is, given any sequence and any position, is there always a constant-sized Mamba block that can execute the COPY operation? We have not provided an answer to this. However, we would like to point out that Mamba blocks of size $O(N)$ always can perform the operation. From this perspective, we state that constant-sized models may have a bottleneck in our paper. However, as pointed out by reviewer PWf8, we also recognize that this explanation is somewhat unclear and could be misleading for readers. We will revise this statement in the subsequent version to make it clearer. Furthermore, as with the theoretical analysis in previous work, our theory focuses on the expressive power of the Mamba block and whether this ability can be learned beyond the scope of this paper is a topic we will explore further in the future.

We greatly appreciate the reviewer's careful review and constructive suggestions. Below are our responses.

W1: The presentation is not good. Notations are not concise enough for readers. One solution is, the authors can provide a notation list and figure out a way to share notations with literature. Figure 1 is a good idea. The authors could integrate equations into the figure. I also feel some texts and equations are redundant in the main texts. For example, preliminaries (introduction to MAMBA) are too long and only Eq 1 is closely related to this paper. The paper should be more dense and concise.

Thank you for your suggestion. We will include a notation list in the subsequent version; we think it's a great idea. We have also noticed that some parts of Section 3 are indeed unnecessary, and we will streamline that section to avoid overloading the readers.

W2: Motivation is not clear. The authors claim that their work focuses on token-level tasks (line 115) in contrast to existing works. But all the discussions and experiments are based on copying a sequence...

Furthermore, it is not clear how copying connects to dynamic programming. Does the authors study in the setting that the model only do copy for DP? Can you try some toke-level copying tasks like LAMBADA.

In Section 4, the overall logic is as follows: for Mamba blocks of constant size, we have identified sufficient conditions for executing the COPY operation, as outlined in Theorem 1 and Line 297. We would like to emphasize that satisfying this sufficient condition is stringent. Of course, even in more relaxed situations where the sufficient conditions are not met, it is still possible for a Mamba block to perform the COPY operation. However, we want to point out that Mamba blocks of size $O(N)$ always can. While we need to provide a sequence $X=[x_i]_{i=1}^{N}$, the theoretical results focus on whether for each $x_i$, the historical record at $pos(i)$ can be copied, rather than copying the entire sequence as a whole.

In Section 5, when solving the DP problem, the model can use the COPY operation to get answers for input or intermediate values to compute subsequent results. Our setup is not "only perform copy for DP," but rather the model solves various operations required in the DP process by combining the COPY operation with those demonstrated in the Lemmas in the appendix.

As for LAMBADA, we directly tested the pretrained Mamba and Transformer models on the LAMBADA test data. For Mamba, we selected the three sizes used in the paper: 360M, 1.4B, and 2.8B. For the Transformer, we used the Pythia transformer models (Biderman et al., 2023) of sizes 410M, 1.4B, and 2.8B, as also adopted by Jalessi et al., 2024 [2]. We found that Mamba outperformed Transformer in terms of accuracy across all three sizes (for example, Mamba with 360M has 53% accuracy while Transformer with 410M has 44%). We suspect that this may be due to (1) we used the pretrained models directly without fine-tuning. Fine-tuning would involve appropriately segmenting the LAMBADA training text to create suitable training data (e.g., ensuring that the last word is meaningful and appears within a certain distance in the preceding text), and we have not yet performed training on LAMBADA's training text; (2) the distance between the target word to be predicted and its previous occurrence is not very large in the test data (the distance ranges from 10 to 80). We plan to further explore and design better approaches for token-level tasks in the future. For example, using tokens instead of a sequence of numbers in tasks like the Phonebook task might be a better choice.

[1] Pythia: A suite for analyzing large language models across training and scaling.

[2] Repeat After Me: Transformers are Better than State Space Models at Copying

We sincerely thank the reviewer for the constructive questions and suggestions. We have noted the reviewer's confusion regarding the circular dependence of definition mentioned in the weakness section. We would like to firstly provide a more detailed discussion of Q1 and Q2.

Q1: the task definition depends on $\delta$, but then the theorem states that $\delta$ observes some constraint; it would be helpful to discuss this circular dependence

You may understand it as follows: Under the premise of Assumption 1, when a Mamba block additionally satisfies Eq. (7), it can perform the COPY operation. In this case, under the constraints of Assumption 1, it naturally follows that $\delta \leq \frac{\epsilon}{(L-1)M||\Delta||\_{\infty}}$​, as demonstrated in the proof. In Assumption 1, we did not impose specific restrictions on the values of $\rho$ and $\delta$, only the relative relationship $\rho \geq \delta$. Under the scenario described in Theorem 1, it further has that "$\rho \geq \text{the right-hand side of Eq. (7)} \geq \delta$" and "the second condition of Assumption 1" $\Rightarrow$ "$\delta \leq \frac{\epsilon}{(L-1)M|| \Delta|| \_{\infty}}$".

Q2: the task definition depends on bi,ci which depends on the Mamba block, but then the theorem shows the existence of a Mamba block given the task. Again, it would be helpful to discuss this circular dependence

The definition of the COPY operation indeed depends on different Mamba blocks: the parameters of different Mamba blocks will yield different attention scores $c_i^T b_j$ for all $i \le N$ and $j \le i$, as well as different $v_i$ for all $i$. However, as long as their output satisfies $o_i \approx v_i$, we consider that they have performed the COPY operation. When the input sequence $X = [x_i]_{i=1}^{N}$ and the position $pos(i)$ we aim to copy is fixed, additionally given some $\delta$ and we consider those Mamba blocks that satisfy Assumption 1, what are the sufficient conditions for these Mamba blocks to achieve the COPY operation? We aim to identify a subspace of the hypothesis space that satisfies these conditions, and this discovery is existential. We suspect that much of the confusion may stem from the role of $\delta$. As explained in the response to Q1, if the initial value of $\delta > \frac{\epsilon}{(L-1)M||\Delta||\_{\infty}}$, then either Assumption 1 or Eq. (7) will not be satisfied.

In addition, we have also noted the reviewer's concerns about the limitations of the theoretical results mentioned in the weakness section. Below is our response.

Weakness part：More disccsuion on Limitations of Theorem 1 & Better writing presentation

For Mamba blocks of constant size, we have identified the sufficient conditions for performing the COPY operation, as stated in Theorem 1 and Line 297. We would like to point out that satisfying these sufficient conditions is stringent because we are considering a relatively general scenario for a given sequence $X = [x_i]_{i=1}^{N}$ and $pos(i)$. Of course, even under less restrictive conditions that do not meet the sufficient criteria, it is still possible for a Mamba block to achieve the COPY operation.
However, the underlying question is: for any given sequence and position, is it always possible for a constant-sized Mamba block to perform the COPY operation? We have not provided a definitive answer. However, we want to emphasize that Mamba blocks of size $O(N)$ can always achieve this.

Our expression is indeed somewhat unclear. In the next version of the paper, we will revise the relevant sections based on your valuable feedback and that of other reviewers. We will reorganize the language used in the definitions, assumptions, and theorems to ensure clearer and more rigorous expression. Additionally, we will provide a more thorough discussion of the limitations of the theoretical results to avoid confusion or misunderstanding for the readers.

I appreciate the clarifications from the authors and the effort they have invested in this work. I think the paper cannot be published as-is (because, judging from confusion it has generated among the reviewers including myself, this work will lead to a lot of reader confusion). But, I believe that, with major revision and improvements on presentation, clarity and framing, this work will become valuable for researchers in relevant domains. I look forward to the next version of this work.