Mimetic Initialization Helps State Space Models Learn to Recall

Thanks for your thoughtful review! We try to address some potential misunderstandings below, and provide a few more experimental results based on your comments.

Weaknesses

1. We'll improve the flow of the paper; the suggestion to move the pretraining section closer to the intro seems reasonable to us, and this makes the connection to previous work on mimetic initialization somewhat more clear.

2. We disagree. Only one of the four components of our initialization is similar to that from Trockman 2023, and it is quite unprecedented that it works on language. This paper showed strong results on vision but marginal results on language, so our positive result on language using a similar technique is quite novel. Setting the parameter $A$ to make state space layers close to linear attention is novel, to our knowledge.

3. This is a significant misunderstanding. Setting $A \approx 1$ is not an "identity prior" -- this controls the receptive field of the learnable causal mask. If anything, this makes the layer behave less like the identity. We make no attempt to avoid vanishing/exploding gradients in the sense of previous work on unitary/orthogonal RNNs. Like previous work on mimetic initialization, we do not consider the "signal propagation" perspective at all, but rather the high-level statistical structure of the weights as inspired by pretrained models.

4. We directly addressed a notable shortcoming of SSMs such as Mamba, which is that they have been shown to be bad at copying/recall [0]. Improving full-scale pretraining results was outside the scope of the paper, which sought to address this one problem with a simple, intuitive solution. The purpose of our paper was primarily scientific rather than application-driven. For example, we contribute to the hypothesis behind mimetic initialization, that the benefit from pretraining comes from storing transferrable knowledge and serving as a good initialization, and that the second component could be isolated in closed-form. We also show that large but constant state sizes endow models with a surprising amount of copy/recall capacity.

I thank the authors for their thorough response. However, many of my questions remain unaddressed and others that were failed to resolve my concerns. I will go over some of the major ones below.

1. This is a significant misunderstanding. Setting $A \approx 1$ is not an "identity prior" -- this controls the receptive field of the learnable causal mask. If anything, this makes the layer behave less like the identity. We make no attempt to avoid vanishing/exploding gradients in the sense of previous work on unitary/orthogonal RNNs. Like previous work on mimetic initialization, we do not consider the "signal propagation" perspective at all, but rather the high-level statistical structure of the weights as inspired by pretrained models.

To clarify, when I said identity prior I was speaking with respect to the eigenvalues of $A$ all being 1 in which case there is no decay across the time dimension. This is precisely what linear attention does, since it has no decay factor, which is why its receptive field does not decay as a function of backwards time. Analogously, the closer the eigenvalues of $A$ lie to 0, the faster the signal along the time axis vanishes. Can the authors please explain to be why the proposed mimetic initialization scheme cannot being viewed as placing a prior on the network to avoid quickly decaying signals (or gradients) as a function of backwards time? I agree, you are not considering the signal propagation perspective, but I am noting a similarity between the two. Please let me know where my logic fails.

2. We concede that GLA may be a reasonable comparison to put in the paper. GLA is quite similar to the state space layer in Mamba 1 and 2, and the same technique from our paper applies: try to set underlying parameters such that the gating parameter is close to all-1s, making the layer resemble linear attention. The correlated C/B (equivalently, query/key) phenomenon also applies. However, in our experiments, we were not able to train networks using only GLA to copy. We're unsure why GLA on its own seems to be unable to learn to generalize on copying, but this makes us think that including it in the paper may be relatively low priority.

This, in my opinion, is not a sufficient rebuttal. The proposed initialization scheme not working on GLA (which the authors have agreed is quite similar to Mamba) raises questions regarding the robustness and generalizability of the approach. Are the authors claiming that mimetic initialization is only effective on Mamba because its initialization scheme is poorly misaligned with linear attention? If so, this is too narrow of a scope to constitute a conference-worthy paper. If the authors were able to show a suite of linear RNN architectures that like Mamba benefit from the linear attention prior, this would be much stronger.

Side note: I believe the reason the GLA architecture aligns with the linear attention prior is due to the normalization factor $\beta = 16$. Reducing the normalization factor would cause it to deviate more from that prior. It is possible that the author's proposed changes to the GLA initializer would affect recall performance in this regime, but this has not been tested

3. Yes, we're using the short conv block before the SSM block. We just replaced the state space layer with linear attention. We'll clarify this in the paper.

Firstly, have the authors ran any experiments without the use short conv block? Many SSM-based architectures outside of Mamba do not use this architectural component. But the claim in the paper is that mimetic initialization helps SSMs recall (not that it helps models with SSM + conv blocks learn to recall). Of course, if the authors are restricting the scope of their claims to Mamba (i.e. when they use the term SSM they are really only referring to Mamba which it seems they may be), then this point is moot.

4. Note that S6 is actually just the Mamba 1 state space layer, so we did study that.

So the authors did study Mamba without the conv layer (i.e. the SSM layer in isolation)? Or are the authors saying that they studied in indirectly via Mamba?

5. Further, if an architecture is in fact the same "class" as Mamba 1/2, then our technique and the accompanying intuition extends very easily. We'll make this clearer in the paper.

I agree that the technique extends readily to most SSMs. My question is not that, but rather if it actually improves performance in other models or if it is specific to Mamba due to Mamba having a "poor" initializer. And if it restricted in scope to Mamba, why does it not improve performance on recall in other models? Is it because initializers in other SSMs already are close to linear attention? Currently, it appears that authors do not have a good understanding of this

Thanks for your review! We try to clarify some things below.

Why do SSMs struggle with copy/recall, how does the init address this?

We meant for this to be clearer in the paper, but it's largely because the learnable causal mask has very small "effective receptive field" at initialization time (the layers only look back into the relatively recent past), which is addressed by our initialization. As we mention in the paper, initialization has historically been important for the success of SSMs in deep learning.

Could the initialization improve language-based tasks?

The point of our initialization was to fix a noted shortcoming of SSMs, that they struggle on copying and recall tasks in particular. Improving non-recall/copying performance is outside the scope of the paper. We suspect that our technique could help at larger scale, especially by using the init for just a fraction of the SSM layers or heads in Mamba 2 -- not applying it to all layers would surely prevent any regression on non-recall tasks.

Broader implications of the work

We tried to highlight the main implications in the introduction: for one, a noted disadvantage of Mamba is that it struggles to learn to copy, and we greatly improve it in this respect through our initialization. It's a vote of confidence that SSMs could be practical for long-range recall tasks, which would greatly speed up inference.

As we mention in the conclusion, our work is related to a hypothesis from previous work on mimetic initialization. Namely, that pretraining can be decomposed into storing transferrable knowledge + serving as a good initialization, and that it may be possible to disentangle the good initialization component. Our paper presents one constructive example of this. Paired with previous work on the topic, we think this contributes to mechanistic understanding of deep learning.

Why is there such a large difference between Mamba 1 and Mamba 2?

First, note that our init often still results in very large gains for Mamba 2 in addition to Mamba 1. The question of why Mamba 2 generally works better has not been resolved, but was addressed by the authors of Mamba. They speculate that the restricted structure of the SSD layer could be helpful [1]. We think ablating all the architectural differences between the two models is outside the scope of our paper.

Attention patterns should be explained more clearly

We explain how we visualize the Mamba attention maps in lines 162-168. The technique is the same as that of [0]. The inputs are the incoming activations $X$, as well as the state space parameters $A, B, C, \Delta$. For softmax self attention, we visualize attention maps in the standard way (we just display $\mathsf{softmax}(QK^T)$) -- there is only one head in our hybrids. We do average over "heads" for the Mamba attention maps as it is cumbersome to look at all of them, but the cross-head variance is quite low. We'll try to clarify this in the paper, please let us know if you have particular suggestions.

How do the maps evolve over time / depend on the seed?

We didn't explicitly study this, though we can say that the results are generally quite robust across trials except for very deep models, where there is more variance (and the initialization greatly increases the proportion of models that successfully learn the task), as in Fig. 3.

We'd assume the Mamba maps align with the self-attention ones fairly rapidly based on the rate of change in accuracy in the plots (the models learn the task rapidly in few steps), which might make this hard to study. We also think it's important to point out that we don't explicitly try to align Mamba layers with self-attention behavior, it's just that those initialized to be closer to linear attention tend to have similar structure, qualitatively, to self-attention maps.

Thanks for your review, we've answered some questions below:

Does this lead to regressions on non-recall tasks?

We think this is unlikely, especially because the technique is effective even when we use it for just one layer (see line 321). Further, pretrained models tend to learn a few SSM layers which implement recall abilities (Fig. 11). Taken together, we think this suggests that initializing just a few layers out of, say, 48 or 64 in a large pretrained model would not harm performance on non-recall abilities.

Is this important with the rise of hybrid architectures?

While hybrids are becoming more popular, full-context self-attention layers remain the most expensive component in those architectures. If we could replace them with a recall-specialized SSM, we would reduce inference costs. Consider a hybrid where the SSM layers are merely specialized rather than using full attention.

No pretraining results

The main purpose of the paper was to understand why SSMs don't learn to copy/recall very well, as in [0]. We effectively diagnose the problem, allowing for better understanding of the capacity of SSMs and their tradeoffs wrt Transformers. The secondary purpose was to investigate the idea of "mimetic initialization", where the effect of pretraining can be viewed as storing transferrable knowledge and serving as a good initialization -- this second effect being something we were able to localize constructively with our technique + analysis of pretrained models. That is to say, the purpose of the paper was largely scientific, rather than attempting to improve production-grade models.

That said, we've tried our technique on a 500M-param Mamba 2 trained on 20B tokens of SlimPajama. We see a small improvement to training and validation perplexity using our technique to initialize a small number of layers. However, we lack enough experimental evidence to endorse this in the paper, and full-scale investigation of our technique for pretraining is outside the intended scope.

I have read the authors response, as well as the other reviews and their corresponding responses. Since multiple other reviewers mentioned the lack of evaluations on language tasks, and the lack of analysis on the effect on pre-training, it still remains a strong concern of mine. While attention layers can be expensive for large context sizes, in most hybrid architectures they are relatively sparse (meaning only a small fraction of the layers are attention layers). Hybrid architectures have demonstrated to be performant and relevant, so I believe my original point still stands. Overall, I do not feel comfortable recommending to accept this paper primarily due to the lack of analysis on pre-training and language tasks. I maintain my score.

We're willing to change the title to reflect the fact that our paper focuses on Mamba SSMs in particular if you think this is a significant improvement. But we'd like to clarify some potential misunderstandings first.

The paper is mostly about the copying task, which previous work has noted state space models struggle on in particular [0]. Consequently, we focus on SSMs which can actually perform on the copying task. This excludes LRU and RetNet: they are linear time invariant and cannot generalize on the copying task. Selective SSMs like Mamba 1, Mamba 2, and GLA implement an input dependent operation that can actually generalize on copying tasks.

Our usage of "State Space Models" is in line with previous work [0,1], assuming we're restricted to SSMs which can copy and recall (beyond one fixed sequence length).

We concede that GLA may be a reasonable comparison to put in the paper. GLA is quite similar to the state space layer in Mamba 1 and 2, and the same technique from our paper applies: try to set underlying parameters such that the gating parameter is close to all-1s, making the layer resemble linear attention. The correlated C/B (equivalently, query/key) phenomenon also applies. However, in our experiments, we were not able to train networks using only GLA to copy. More experiment details below:

In particular, we replaced the SSD layer inside Mamba 2 with GLA with the same hidden dimension and number of heads. We attempted to set $\alpha_t = \sigma(x_t W_{\alpha 1} W_{\alpha 2} + b_\alpha )$ so that $G_t = \alpha_t^T 1 \approx 1$. In the GLA implementation, $\sigma(x) = \exp( \mathsf{LogSigmoid}(x)/16 )$, so we set $W_{\alpha 1} W_{\alpha 2} \approx 0$ and $b_\alpha \approx 10$ to get $\alpha_t \approx 1$. We compare this to the default initialization, training on 50-length strings, using 4-layer models with hidden dim 1024.

Models using only GLA were not able to learn to generalize on the copying task despite trying several initialization hyperparameters. Our initialization seems to help fit to the training length, at least. We also tried Mamba 2 models where we replaced SSD in only the second layer by GLA (with/without our init), and this performs comparably to Mamba 2 with our init. We note that GLA is actually closer to all-1s using its default init than Mamba 2 due to the particular sigmoid parameterization, so we might expect (at least this component) of our init to have less of an effect.

We're unsure why GLA on its own seems to be unable to learn to generalize on copying, but this makes us think that including it in the paper may be relatively low priority.

Re: "SSM layers cannot mimic full Attention". There's a misunderstanding here; selective SSMs such as Mamba 2 can in fact mimic (linear) attention, and this is in fact noted in the title of the Mamba 2 paper "Transformers are SSMs" [1]. We show this in our paper mathematically. Of course, due to the limited state size and lack of softmax, we can't mimic full (softmax) attention.

We can add some references, though we already cite [3] and the time invariant SSM and linear attention papers are not too relevant (as we use no specialized linear attention techniques). [5,7,8] are good references to add though. Let us know if you have further thoughts.