Stuffed Mamba: State Collapse and State Capacity of RNN-Based Long-Context Modeling

Thank you for your thoughtful review. Regarding the mentioned weaknesses, we make the following response.

The fundamental limitation of constant-size memory.

The constant size of the recurrent state is indeed a fundamental bottleneck of RNN-based language models, and we are not claiming that it is the correct way ahead. This paper represents one of the first steps toward long-context modeling with recurrent language models. The constant-size memory is exactly the reason why we conducted the experiments described in Section 5, i.e., to estimate the training length at which the model can generalize and the context length at which the model can have near-perfect passkey retrieval accuracy (a very simple form of associative recall in natural language text). The result can provide important clues for choosing a suitable context length given the state’s constant size.

Additionally, although the constant size of memory limits the capacity of RNN-based models, it also makes them much more efficient. Therefore, we believe the choice between RNNs and transformers is a trade-off between performance and efficiency, and perhaps some future architecture may be able to combine the best of both worlds. Our work is not meant to propose such architectures, but to analyze how and why RNNs exhibit length generalization failure.

Evaluation of passkey retrieval and perplexity may not reflect real-world tasks.

We have based our analysis on passkey retrieval and perplexity because we believe they are the minimum bar for memory recall. Failure in such tasks means that the model will likely fail on more realistic tasks. The main contribution of the paper is not to show that RNN-based language models are capable of handling realistic long-context problems, but that they are capable of performing the simplest kind of contextual information retrieval with lengthy contexts beyond 100K tokens. This is an ability that we have not observed in prior RNN-based models without specialized training.

Rigorous definition of state capacity

Thank you for bringing this up. The current draft of the paper has conflated two different notions of state capacity: the theoretical capacity of the state and the empirically estimated capacity of the state. In short, the "state capacity" mentioned in Figure 11, Section 5, and Section 7.2 all refer to the empirically estimated state capacity (by sweeping different training lengths). The measurement of state collapse is explained in Line 440. While a rigorous theoretical analysis is interesting, we have not included such an analysis in this work because our goal is to find the relationship between this estimated capacity and the state size. In Figure 11, we empirically find that this relationship is linear, which is a novel discovery. This relationship is an important consideration for RNN-based model training to avoid state collapse. An analysis of the theoretical state capacity does not change this conclusion, hence, we have chosen to leave that for future work.

We will kindly notify you with another comment when we have updated the PDF file, and we will update all the reference to the paper in this response as well.

Update: We have now revised the paper and updated the PDF file.

Thank you for the thoughtful review.

Before we respond to each of the mentioned weaknesses, we would like to emphasize the priorities of the paper because we believe that you have focused on parts that we believe are less important. The main contribution of the paper is:

1. The discovery and systematic analysis of state collapse (SC) (Section 4.1 and 4.2)
2. We present the state overparameterization hypothesis as an explanation for SC. This establishes a relationship between SC and the state capacity. (Section 4.3)
3. We empirically discover that there exists a minimum training length such that training lengths beyond this threshold result in a model without SC. Thus, this strengthens the hypothesis about state overparameterization (Section 7.2, Figure 9).
4. We empirically discover that this training length threshold is a linear function of the state size and that the model can recall contextual information beyond this training length (Section 7.2, Figure 11)

The training-free mitigation methods (Section 4.4.2) are meant to showcase/strengthen the analysis conclusions we have made in Section 4.1 and 4.2, which is why we put it in a \subsubsection. In the updated version of the paper, we will revise the abstract and introduction to reduce the amount of content related to the training-free mitigating methods. We will move the evaluation result of these methods to the appendix because the paper will be more readable by using more space to explain our reasoning behind the state overparameterization hypothesis and empirical findings regarding the state capacity. Moreover, we will remove Method 2 (State Normalization) from the paper altogether, because the performance is bad and it is not the focus of our paper.

In response to your summarization of weaknesses, we make the following remarks.

Comparison between the mitigation methods. Due to limited space, we have not compared the performance of the mitigation methods, because they are less important, and the main take-away of the paper is the analysis of SC, the cause of it, and the training-based result.

Generalizability of mitigation methods. Once again, we do not consider the generalizability of these methods because they are less important, and we have decided to spend more time and space on the analysis and the training-based experiment.

Trade-offs at shorter sequence lengths of mitigation methods. This is an existing limitation, but since these methods are not the most important part of our paper, we leave such an exploration for future works.

Engaging more relevant prior works. We will add the mentioned papers and discuss how they are relevant to our paper in the updated draft. Thank you for the suggestion.

Visualization of the results. We will update the figures to make them clearer. The term “Ans. Depth” is used in existing works as well, so we thought it would be easily understandable. The color scheme will be added.

Questions

• How general are these methods? Regarding the training-free mitigation methods, as mentioned in Section 4.4.2, Method 3 is generalizable to all variants of RNN that can be written as a weighted sum of token representations, which include RWKV 4, 5, and 6, GLA, HGRN-2, RetNet, and many other contemporary RNNs. As for Method 1 and 2, it should be fairly easy to apply them to other RNN architectures, but that is outside the scope of this work.
• Of the proposed methods, which are more or less effective? Our preliminary experiments show that Method 3 is the most performant in most cases, but we leave a more comprehensive comparison for future work.
• Is the tradeoff at short sequences inherent? Method 1 and 2 will weaken the inserted information, therefore, it likely results in a performance drop across different context lengths. For method 3, the trade-off at shorter sequences can be entirely avoided by tweaking the sliding window size. We believe that the trade-off at short sequences is not inherent, but a more thorough exploration is outside the scope of our work.

Update: We have now revised the paper and updated the PDF file.

Thank you for the detailed and constructive review.

Regarding the precise definition of SC.

The precise definition of State Collapse (SC) is as follows:

SC refers to the exploding values in the recurrent state that causes performance degradation.

Moreover, with careful consideration, we have decided to rename SC to State Explosion (SE) in the updated paper, which we believe will make it clearer that the term refers to the value of the state and not the generalization failure. In the response below, we will use the new terms SE and SC interchangeably.

In Section 4.2, when we say the "cause of SC", we refer to the cause of these exploding values, which is unexpected because the recurrent formula is stable as mentioned in Line 250. We acknowledge that the title of Section 4.2 is misleading, we are not claiming that the exploding variance is the cause of SE, this section is meant to illustrate the distribution of the state to better understand the phenomenon. In the updated paper, we will rename Section 4.2 to "Observation of State Explosion", and we will provide the above precise definition of SE in that section.

Explanation of Line 329 and Line 400.

The Section reference number is mistakenly incorrect, we are supposed to refer to Section 7.2. It is an empirical observation that we discover by training on different training lengths, and checking if SE is exhibited on the "newlines" prompt (as explained in Section 5). In Line 351, the claim about "if and only if" is an hypothesis, and is meant to be confirmed by the experiments in Section 7.2. We will rewrite the mentioned sections to make this clearer in the updated paper.

The Definition of State Capacity

We acknowledge that we have conflated two different concepts: the theoretical capacity of the recurrent state, and the empirically estimated state capacity. Thank you for bringing this up. In short, the "state capacity" mentioned in relation to state overparameterization, Line 408, Section 5, and Section 7.2 all refer to the empirically and indirectly estimated state capacity (by sweeping different training lengths). We will explicitly highlight this distinction in the updated paper (in Section 5).

Missing Related Works

Thank you for mentioning these papers. We will include all of them in the Related Works section in the updated paper and explain how they relate to and differ from our paper.

Other Notions of Capacity

Thank you for bringing this up. We want to first emphasize that the experiment regarding language modeling (Figure 11, left) is meant to show the empirical linear relationship between the training length threshold (beyond which SE is not exhibited) and the state size. This linearity between the training length threshold and the state size is a novel discovery, to the best of our knowledge.

Regarding passkey retrieval, we chose to search for the empirical memory capacity on this task because we believe it can serve as the minimum standard for working associative recall. Namely, the failure in passkey retrieval will likely imply the failure in more challenging tasks such as multi-query retrieval and benchmarks such as RULER. We will dedicate a part of the updated paper to discussing the reason for choosing these tasks and how they are related to other tasks.

Questions 1 - 3

We believe these questions are addressed in the above comment. Thank you again for mentioning these points.

Question 4

For the SE mitigation method 1, the intervention is done by scaling $B_t$ and $\alpha_t$ with a multiplier. The multiplier is chosen by sweeping different values and validating on the "newlines" prompt. The final values are given in the legend in Figure 9.

We will kindly notify you with another comment when we have updated the PDF file, and we will update all the references to the paper in this response as well.

Update: We have now revised the paper and updated the PDF file.

Thank you for your detailed review.

Definition of state collapse (SC)

Regarding the definition of State Collapse (SC), we acknowledge that we have mistakenly left out a precise definition in the paper. After a discussion between the authors, we have decided to make the term more precise by changing it to "State explosion" (SE), which emphasizes that the term describes an observed change in value. In the following response, we will use SE instead of SC. In the updated paper, we will add the following definition to Section 4: SE is the phenomenon where the recurrent state exhibits exploding values, causing performance degradation.

Difference between SE and length generalization failure

Thank you for raising this question. SE is one possible cause of length generalization failure, but it is not the same as length generalization failure. The reasons we hypothesize that SE is caused by overparameterization are the following empirical observations:

1. We observe that the performance drop is much more severe for models with larger state sizes.
2. We observe that during training, the performance within the training length consistently increases with the training amount, but the performance beyond the training length will first increase and then decrease after a certain amount of training. This implies that the model can overfit the training length, which is an indication of overparameterization. Since the contextual information is only stored within the recurrent state, it is reasonable to infer that the state is overparameterized.
3. The first token is not forgotten well beyond the training length, which indicates that the model is trying to retain all contextual information. This is only possible when the state is sufficiently large because, with a sufficiently small state (e.g., a 2x2 state), it will be impossible for the model to retrieve information with high precision from the state containing information about 8K tokens. Hence, we conclude that, without an overparameterized state, the model will learn to generate a sufficiently small memory decay value $\alpha_t$ (i.e., it induces more forgetting), which can help avoid the exploding values. The connection between smaller memory decay value and SE is highlighted in Figure 7, where the heads that exhibit SE have larger memory decay value (i.e., less forgetting). This strong correlation between larger memory decay values and exploded heads is consistently observed in different layers and model sizes.
4. Finally, and most importantly, we empirically observe that when sweeping different training lengths (Section 7.2), for every state size, there exists a training length threshold, beyond which the model does not exhibit SE.

Therefore, when we claim that the state is overparameterized, we are claiming that the state is too large and is not sufficiently utilized for the given training length. This causes the model to generate memory decay values ($\alpha_t$) that are too large, which causes the state to explode due to the additive nature of the update rule.

The updated version of the paper will change the content of Section 4.3 to highlight the above reasoning behind the overparameterization hypothesis.

Explanation of LongMamba and the reason to exclude another baseline.

LongMamba is a baseline and not an implementation of mitigation method 1. We do not consider [1] as a baseline because, as stated in Appendix C, it requires task-specific tuning, which is an entirely different application setting. Due to the limited space, we have kept the description of the mitigation method 1 short. In brief, method 1 simply scales $B_t$ and $\alpha_t$ by multiplying them with a scalar. We choose the two locally optimal sets of multipliers by validation on the newlines prompt. The legend in Figure 19 for Method 1 shows the different multipliers we arrived at. We will revise the paper to include this explanation.

Finally, we would like to emphasize that the training-free mitigation methods are not the main focus of the paper, and kindly ask the reviewer to focus on the experiments/findings related to SE and state capacity when evaluating the paper. To make the priorities of the paper clearer, we have decided to revise the paper by moving the evaluation result of the training-free mitigation methods to the appendix, removing Method 2 (State Normalization) altogether, and using the extra space to provide an additional explanation of the findings in Section 5 and to explain the reasoning behind our state overparameterization hypothesis as mentioned above. We will also revise the abstract and introduction to mention less about the training-free mitigation methods

We will kindly notify you with another comment when we have updated the PDF file, and we will update all the references to the paper in this response as well.

Update: We have now revised the paper and updated the PDF file.

(cont.)

2. Limitations with training-free mitigation methods (Claim 2).

In response to this weakness, we would like to first emphasize that the training-free mitigation methods are not the focus of this paper. They should only serve as an empirical demonstration of the insight that SE can be avoided by inducing forgetting (based on the analysis in Section 4.1 and 4.2). As a result, we have only dedicated one page of space to these methods. In practice, practitioners are encouraged to rely on training-based methods to mitigate SC (as described in Section 4.4). Training on longer sequences is not too expensive and can yield much better results. We kindly ask the review to focus on the experiments/findings related to SE and state capacity. To further emphasize the focus of this paper, we will revise the abstract and introduction to reduce the amount of content related to the training-free mitigation methods, and we will move the evaluation result of these methods to the appendix. We will also remove state normalization from the paper altogether because it performs poorly.

• Methods 1 and 3 induce forgetting. Although these methods induce forgetting, they still outperform the base model with SC. This is a trade-off for the benefit of training-free nature. To mitigate SC without inducing forgetting, practitioners should rely on the training-based method.
• Perplexity still rises on longer sequences with methods 1 and 3. This is an existing limitation of our method, however, this is already much better than the base model and the existing method.
• Method 2 makes training/prefill non-parallelizable. Since the mitigation method is applied post-training, this only affects the prefilling phase, which is still a significant limitation. However, the same limitation applies to any non-linear RNN, and optimizing the efficiency of such architectures is an open problem that we do not tackle in this paper.
• Method 2 degrades performance substantially. Yes, the model is not able to adapt to state normalization without training. However, with training, we believe it is better to simply use a much longer training length, as we did in Section 7.2.

3. Evaluation on a limited set of tasks (Claims 2 + 3).

Once again, the training-free mitigation methods are meant to strengthen/showcase our analysis conclusions, therefore, we did not evaluate them more thoroughly in the paper. We will evaluate these methods of passkey retrieval in the updated paper. In short, Method 1 and 2 have very poor retrieval performance, Method 3 can only correctly retrieve the passkey when it is located inside the window. For the training-based method, we will add an evaluation of RedPajama in the Appendix. Regarding other benchmarks, the models we experiment with are too small to have good results on more challenging tasks other than passkey retrieval.

Questions

• As stated, the analysis can be performed on any layer with heads whose state collapses. We arbitrarily chose one of the layers.
• We have changed Figure 19 into the log-scale in the updated paper.
• We have added the color scale in Figure 10. In brief, the value in each cell is very close to 1.0.

We have now revised the paper and updated the PDF file.

Dear reviewer, thank you for your detailed and thoughtful review. In response to your review, we would like to make the following remarks.

1 Evidence for state over-parameterization (Claim 1) is unclear and limited.

It’s unclear how the layer/heads were chosen. Furthermore, the causal link between high memory strength in this one layer and state collapse is not evident.

The statistics of the state for every layer are already presented in Figure 17 in Appendix G. In short, several layers exhibited SC, and we arbitrarily chose one of the layers (layer 38) as an example due to limited space. The exact same analysis and findings can be made on other layers with SC as well (e.g., layer 9, 11, 17, 26, 30, etc). Similarly, we chose to only show the statistics on the first 8 heads to keep the plot simple, the same findings can also be made on other heads. Additionally, we emphasize that the correlation between high memory strength and SC is consistent in every layer. I.e., for any given layer, heads with SC have high memory strength and heads without SC have low memory strength.

Looking at Figure 6 it’s not clear that $B_t$ "explodes" $\Delta_t$ before. It also isn’t obvious why this would mean that the collapse is "largely attributable to $B_t$".

Thank you for this notice. Our claim about $B_t$ exploding before $\Delta_t$ is limited to the heads with SC. In Figure 6 (a), the curves correspond to the 8 heads shown in Figure 4, therefore, only the head 2, 4, and 7 (green, purple, and gray, respectively) are collapsing. Therefore, Figure 6 (a) shows that the explosion of $\Delta_t$ is only observed after 20K tokens, while in (b), we see that $B_t$ already shows considerable signs of explosion before 20K (the values are considerably larger than the one within the training length). However, we acknowledge that that claim about SC being largely attributable to $B_t$ is too bold and it will be removed in the updated paper.

Lack of support in Section 4.3.2

This is a mistake, this sentence is meant to refer to Section 7.2, in which we swept different training lengths, and empirically found that for each state size, there exists such a training length threshold beyond which SC is not exhibited.

Confusion about the state overparameterization argument. Overfitting seems unrelated to the issue at hand. Training data in the language modeling setting is abundant, so it doesn’t seem like overfitting is part of the problem?

The definition from Wikipedia is "An overfitted model is a mathematical model that contains more parameters than can be justified by the data."[1]. We are claiming that the state overfitted the training length because it contains too many parameters, resulting in an update rule (in the heads with SC) that fails to forget past information to avoid exploding values. The abundance of training data does not prevent this overfitting because, for the given training length (e.g., 8K), the distribution of the state seen during training is not varied enough to justify the state parameter count. This is a hypothesis we have made based on the inspection of the activations (Section 4.2), and then strengthened by Section 4.4, the mitigation methods, and the observation that there exists a training length threshold beyond which SC does not exhibit.

References: [1] https://en.wikipedia.org/wiki/Overfitting