Mimetic Initialization Helps State Space Models Learn to Recall

## Abstract

Abstract Recent work has shown that state space models such as Mamba are significantly worse than Transformers on recall-based tasks due to the fact that their state size is constant with respect to their input sequence length.
But in practice, state space models have fairly large state sizes, and we conjecture that they should be able to perform much better at these tasks than previously reported.
We investigate whether their poor copying and recall performance could be due in part to training difficulties rather than fundamental capacity constraints.
Based on observations of their “attention” maps, we propose a structured initialization technique that allows state space layers to more readily mimic self-attention. Across a variety of architecture settings, our initialization makes it substantially easier for Mamba to learn to copy and do associative recall from scratch.

## 1 Introduction

State Space Models (SSMs) show promise as a potential replacement for Transformers (Vaswani, )
with substantially lower inference costs (Gu & Dao, ; Dao & Gu, ).
While Transformer memory grows linearly with the input sequence length,
SSMs use only a constant amount, compressing all the context into a fixed-size state.
SSMs perform comparably to Transformers on a variety of common benchmarks.
However, recent research has highlighted a set of tasks on which SSMs perform substantially worse than Transformers (Waleffe et al., ),
particularly those involving copying or recall (Jelassi et al., ; Arora et al., ).
This is perhaps unsurprising, as it is harder to recall from a compressed, fixed-size representation, particularly as its length grows.

Nevertheless, SSMs use relatively large state sizes in practice, and we wonder if their
poor performance on tasks such as copying could be due to training difficulties
rather than fundamental capacity constraints.
We present a qualitative study of the failure modes of SSMs on the copying task.
In particular, we inspect the time-dependent linear transformation matrix of Mamba layers,
which is analogous to the attention map of self-attention layers.
We compare these layers to their counterparts in self-attention/Mamba
hybrid architectures that successfully learn to copy,
and based on these comparisons,
we propose a structured initialization
technique that allows Mamba layers to more readily mimic self-attention.
Our technique makes use of the fact that state space layers
can be seen as a form of linear attention with a learnable, structured
causal mask.
We find evidence that such linear-attention-like Mamba layers arise
naturally after large-scale pretraining, suggesting that this
pattern may be fundamental to the recall abilities of SSMs.

The proposed mimetic initialization allows Mamba to quickly learn
to copy and do associative recall on up to
$\mathbf{4\times}$ longer strings, and we show for the first
time that SSMs can achieve $\mathbf{2\times}$ length generalization or more.
Mimetic initialization is essentially compute-free,
but we show it is comparable to pretraining
in allowing Mamba to learn to copy and recall.
Our work helps to better understand
the capacity of SSMs relative to Transformers in practice
and can assist in further studies of their capabilities,
which may have been underestimated by previous research.

### Related work

Recently, Jelassi et al. () did a thorough investigation of the ability
of state space models (in particular Mamba 1) to copy in comparison to Transformers.
Their theoretical results demonstrate that SSMs with a fixed state size have fundamentally
limited copying capacity,
unlike Transformers which can strongly generalize.
Empirically, they find that Transformers (especially with their
proposed custom position embeddings) vastly outperform
SSMs on copying, both in terms of learning and length generalization.
They note that in practice, SSMs may be better at copying than expected
due to their relatively large state sizes, but do not observe very good
copying performance in their experiments.
Similarly, Arora et al. () note that SSMs struggle on recall tasks
due to their limited state size.
They propose an effective
intervention in the form of interleaved kernelized linear attention layers
that boost recall performance.
The second, improved version of the Mamba architecture improves upon
associative recall ability, although the authors note that this task remains
difficult for SSMs (Dao & Gu, ).

Initialization has been important for SSMs since
their introduction to deep sequence modeling by Gu et al. ();
a structured initialization of the state matrix was crucial
to the performance of these earlier time-invariant SSMs (Gu et al., ; Gupta et al., ; Gu et al., ; Smith et al., ).
Our work further demonstrates the importance of initialization for SSMs,
taking inspiration from *mimetic initialization* (Trockman & Kolter, ; Trockman et al., ),
which uses pretrained models as case studies of good initialization.
For example, previous work noted that self-attention layers in pretrained
Vision Transformers may try to imitate the local mixing ability of convolutions,
which is reflected in the correlations between query/key and value/projection weights;
initializing weights with statistical structure that mimics this pattern
greatly improved trainability.
We follow a similar methodology to propose a novel mimetic initialization technique for state space layers
based on our observations that (1) these layers can represent linear attention, which
can improve recall and (2) they sometimes approximate linear attention in pretrained models.

Figure: (a) Training a Mamba with default initialization to copy.
Refer to caption: x1.png

## 2 Preliminaries

Recently, state space models have become popular as a choice
of token mixing layer, i.e., as a replacement for self-attention.
We refer to layers that use state space models for this purpose as *state space layers*.
As it is common in the literature, with a slight abuse of definitions, we refer to architectures like Mamba 1 & 2 that use state space layers only for sequence mixing as *state space models*.

### State space models

For a scalar sequence $x\in\mathbb{R}^{\mathtt{T}}$,
SSMs are linear recurrences of the form

$$ $h_{t+1}=\bar{A}h_{t}+\bar{B}x_{t},\quad y_{t}=Ch_{t},$ (1) $$

where $h_{t}\in\mathbb{R}^{\mathtt{N}}$ is a hidden state,
and $\bar{A}\in\mathbb{R}^{\mathtt{N}\times\mathtt{N}}$,
$\bar{B}\in\mathbb{R}^{\mathtt{N}\times 1}$,
$C\in\mathbb{R}^{1\times\mathtt{N}}$
are the state space model parameters.
Traditionally, SSMs are continuous systems, and the bar notation
refers to the *discretized* form of parameters $A$ and $B$,
which depend on the step size $\Delta$ that is used
to sample an implicit underlying continuous signal $x_{t}=x(\Delta t)$.
Typically,
some structure is imposed on $A\in\mathbb{R}^{\mathtt{N}\times\mathtt{N}}$,
such as diagonal-plus-low-rank (S4), diagonal (Mamba),
or scalar-times-identity (Mamba 2).

In contrast, *selective* SSMs such as the S6 layer in Mamba allow
the parameters $\bar{A}_{t},\bar{B}_{t},C_{t}$ to vary with time,
i.e., depend on $x_{t}$.
The particular state space layer in Mamba operates on sequences of $\mathtt{D}$-dimensional tokens
$X\in\mathbb{R}^{\mathtt{D}\times\mathtt{T}}$. Indexing tokens with $t$ and *channels* with d, it computes

$$ $h_{(t+1),d}=\bar{A}_{td}h_{td}+\bar{B}_{td}X_{td},\quad y_{td}=C_{t}h_{td},$ (2) $$

where $\bar{A}_{td},\bar{B}_{td},C_{t}$
depend on *all* channels of input $x_{t}$, but with different discretization parameters $\Delta_{td}$, hence the dependence of $\bar{A}_{td}$ and $\bar{B}_{td}$ on $d$.
Define the underlying
parameters $W_{B},W_{C}\in\mathbb{R}^{\mathtt{N}\times\mathtt{D}}$,
and $A\in\mathbb{R}^{\mathtt{D}\times\mathtt{N}}$.
Let $W_{\Delta}\in\mathbb{R}^{\mathtt{D}\times\mathtt{D}}$ be a rank-$r$ matrix,
and bias $b_{\Delta}\in\mathbb{R}^{\mathtt{D}}$.
Then the continuous state space model parameters are computed as $B_{t}=W_{B}^{T}X_{:,t}$ and $C_{t}=W_{C}^{T}X_{:,t}$.
The parameters of the discretized state space models are then computed as follows:

$$ $\Delta_{t,d}=\mathsf{softplus}(W_{\Delta d}^{T}X_{:,t}+b_{\Delta,d}),\quad\bar {A}_{td}=\exp(A_{d}\Delta_{t,d}),\quad\bar{B}_{td}=B_{t}\Delta_{t,d}.$ (3) $$

Please refer to Dao & Gu () for a more detailed discussion on selective SSMs.

Figure: Figure 2: A hybrid Mamba architecture with one Self-Attention layer easily learns to copy. Dotted lines: performance on training length (50), solid: $2\times$ length generalization (100).
Refer to caption: x3.png

### Matrix form of SSMs

The operations of Eq. 3 can be written concisely in matrix form:

$$ $\displaystyle\Delta$ $\displaystyle\coloneqq\mathsf{softplus}\left(W_{\Delta}X+b_{\Delta}\right)$ $\displaystyle\in\mathbb{R}^{\mathtt{D}\times\mathtt{T}}$ (4) $\displaystyle\bar{B}_{d}$ $\displaystyle\coloneqq W_{B}X\odot\mathbf{1}_{n}\Delta_{d}$ $\displaystyle\in\mathbb{R}^{\mathtt{N}\times\mathtt{T}}$ (5) $\displaystyle C$ $\displaystyle\coloneqq W_{C}X$ $\displaystyle\in\mathbb{R}^{\mathtt{N}\times\mathtt{T}}$ (6) $\displaystyle\bar{A}_{d}$ $\displaystyle\coloneqq\exp\left(A_{d}^{T}\Delta_{d}\right)$ $\displaystyle\in\mathbb{R}^{\mathtt{N}\times\mathtt{T}}$ (7) $$

As noted first by Ali et al. (), the time-varying discrete recurrence $h_{t+1}=\bar{A}_{t}h_{t}+\bar{B}_{t}x_{t},\ y_{t}=Ch_{t}$ can be unrolled and viewed as a matrix operation.
Namely, channel $d$ of the output of an SSM layer, denoted with $Y_{d}\in\mathbb{R}^{\mathtt{T}}$, can be written as $Y_{d}\coloneqq M_{d}X$, where $M_{d}\in\mathbb{R}^{\mathtt{T}\times\mathtt{T}}$ is a matrix transformation dependent on $d$.
Each matrix $M_{d}$ represents a time-dependent linear transformation, much like
attention maps in self-attention.
For $i,j\in[\mathtt{T}]$, the $M_{d}$ matrix of the Mamba state space layer for channel $d$ can be expressed as follows, where $\mathbf{1}\{i\leq j\}$
does causal masking:

$$ $\displaystyle M_{d,i,j}=C_{:,i}^{T}\left(\Pi_{k=j+1}^{i}\mathsf{diag}(\bar{A}_ {d,:,k})\right)\bar{B}_{d,:,j}\times\mathbf{1}\{i\leq j\}.$ (8) $$

Eq. 8 can be viewed as a
linear attention matrix computed from $\bar{B}$ and $C$
with a learnable causal mask parameterized by $\bar{A}$ (Dao & Gu, ).
As it will be useful later, we note
that in practice, $A$ is
parameterized as $A\coloneqq-\exp(A_{\mathsf{log}})$ with $A_{\mathsf{log}}\in\mathbb{R}^{\mathtt{D}\times\mathtt{N}}$.
The selective state space layer of Mamba 2 is broadly similar to that of Mamba 1;
it follows equations 4–7, but instead of having $\mathtt{D}$ different $A_{d}$ and $\Delta_{d}$, it has $\mathtt{H}$ independent $A$ and $\Delta$, each of which are repeated $\mathtt{D}/\mathtt{H}$ times to construct $A_{d}$ and $\Delta_{d}$.
Each of these $\mathtt{H}$ independent $A$ are parameterized as scalar-times-identity matrices, resulting in just $H$ parameters.
These $\mathtt{H}$ components correspond to “heads”,
leading to only $\mathtt{H}$ unique $\bar{A}_{d}$ and $\bar{B}_{d}$ parameters, and only $\mathtt{H}$ “attention matrices” $M_{d}$ (*c.f.* Eq. 8),
as in multi-head attention.

### Mamba architecture

Mamba 1 and 2 are prominent sequence modeling architectures that combine selective state space layers (as the sequence mixer) with more standard layers.
We describe below the Mamba 1 block, and refer the reader to (Dao & Gu, ) for details on Mamba 2, which are not essential to our work.
Omitting the final LayerNorm, the Mamba block is a composition of two sequence mixer layers (1D convolution and a selective SSM layer) a gated linear block:

$$ $\displaystyle W_{3}\{\mathtt{SSM}\left[\sigma(\mathsf{DepthwiseConv1d}(W_{1}X) )\right]\odot\sigma(W_{2}X)\}+X,$ (9) $$

where $\sigma$ is SiLU (Hendrycks & Gimpel, ).
Mamba 2 simplifies this block, merging all projections into $W_{1}$.
For both, the convolution layer before the
SSM will be considered in our initialization.

### Mamba attention maps

Throughout this work, we visually inspect $M_{d}$
to better understand the operation implemented by Mamba layers.
However, it is infeasible to look at all $\mathtt{D}$ maps, and we instead
visualize and report the average over channels
$\tfrac{1}{\mathtt{D}}\sum_{d=1}^{\mathtt{D}}M_{d}$, which we hereafter refer to as the *attention map* of a Mamba layer.
In practice, the inter-channel variation in maps is relatively
small, as the behavior of $M_{d}$ is dominated
by $\bar{B}_{d}$ and $C$.
We also sometimes find it useful to inspect
the *average attention mask*
$\tfrac{1}{\mathtt{D}\mathtt{N}}\sum_{d=1}^{\mathtt{D}}\sum_{n=1}^{\mathtt{N}}(
\Pi_{k=j+1}^{i}\mathsf{diag}(\bar{A}_{d,:,k}))_{n,n}$
to approximately determine the effective receptive field
of the Mamba layer (i.e., how far into the past it can look).

### Copying task

Most of our experiments focus on copying, a simple
task where SSMs are known to fall far behind Transformers.
We train the model to predict the paste string given the copy string,
emitting a stop token at completion.

$$ $\underbrace{\mathtt{abcdefghijk}}_{\mathsf{copy\ string}}\quad\mathtt{|}\;\;\; \underbrace{\mathtt{abcde}\,\underline{\mathtt{?}}}_{\mathsf{paste\ string}}\cdots\square$ (10) $$

Since Transformers cache the whole sequence,
it is easy for them to learn the task and to generalize
far beyond the training length.
However, since SSMs compress tokens into a fixed-size state,
it is hard for them to store and decode back long sequences.
We consider copying sequences of varying length and of different vocabulary size, drawing tokens uniformly at random.
We also investigate *stack-order* copying,
where the paste string needs to be generated in the reverse order.

### Multi-query associative recall

Another synthetic task that has been shown to be an important
discriminator between Transformer and SSM abilities
is multi-query associative recall,
which tests models’ ability to store and recall
many key-value pairs.
Transformers are well-suited for this task, as they can implement induction heads easily (Olsson et al., ).

$$ $\underbrace{\mathtt{a1\ b2\ c3\ d4}}_{\mathsf{key-value\ pairs}}\quad\mathtt{| }\;\;\;\underbrace{\mathtt{c3\ b}\,\underline{\mathtt{?}}}_{\mathsf{queries}}\cdots\square$ (11) $$

Similarly to copying, we investigate length generalization
on multi-query associative recall.
In our implementation, each key may occur only once,
i.e., it cannot be overwritten by later key/value pairs.

## 3 Initializing state space layers to be more like attention

To better understand why Mamba often fails to learn to copy, we start by examining a small model
trained to copy 50-character strings.
In Figure 1(a), we can see that Mamba plateaus.
Visual inspection of its attention maps reveals that it has probably failed to learn
an interpretable copying operation.

### Attention enables copying

To explore what Mamba might be missing to allow it to copy,
we trained a hybrid eight-layer Mamba whose fourth layer is single-head self-attention.
As shown in Fig. 2,
this one layer enables perfect copying performance, both
on in-distribution length-50 strings (dotted lines) and generalizing
to length-100 strings (solid lines).
The softmax attention head learns a sharp “look-behind” operation,
constructing the paste string by directly attending to the copy string,
likely exploiting an implicit position embedding learned by the preceding Mamba layers.
We propose two initialization changes that allow state space layers
to better use their state capacity.

### 1. State space layers can be linear attention

While there is likely more than one way to learn to copy,
we suspected that Mamba’s copying ability is tied to its ability to represent
a similar operation to the one in this self-attention layer.
Notably, in Figure 1(a), the Mamba layers tend to look only into the recent
past, while the self-attention layer in Figure 2 can attend all the way
to the beginning of the string.
While SSMs cannot look arbitrarily far into the past because of their fixed state size,
even in the simplest time-invariant SSMs,
the amount of history stored in the state is controlled by the parameter $A$,
whose initialization was crucial to the initial success of these models (Gu et al., ).

Consequently, we focus on the state matrix $A$, which controls the
“receptive field” of the state space layer.
Note in Eq. 8 that if $\bar{A}_{d}\approx 1$,
then $M_{d,i,j}\approx C_{:,i}^{T}\bar{B}_{d,:,j}$.
That is, the state space layer’s attention map resembles a product of *queries* and *keys*.
The only inter-channel variation in this equation is
from $\Delta_{d}$ in Eq. 5,
so that if $\Delta_{d}\approx 1$ then
$\bar{B}_{d}\approx W_{B}X$, which results
in $M=X^{T}W_{C}^{T}W_{B}X$, which is simple linear attention before applying
the causal mask.
Thus, if we set parameters so that
$\Delta_{d},\bar{A}_{d}=1$,
the state space transformation is the same for every channel,
and it is simple (non-kernelized) linear attention
with head dimension $\mathtt{N}$ and no value/projection matrices:

$$ $\Delta_{d},\bar{A}_{d}\approx 1\quad\Longrightarrow\quad Y\approx X\cdot \mathsf{CausalMask}\left(X^{T}W_{C}^{T}W_{B}X\right)\in\mathbb{R}^{\mathtt{D} \times\mathtt{T}}.$ (12) $$

However, both $\bar{A}_{d}$ and $\Delta_{d}$
are parameterized and input-dependent, so we cannot directly set them to one.
We use details of the Mamba implementation: To make $\bar{A}_{d}=\exp(A_{d}^{T}\Delta_{d})\approx 1$, we parameterize $A=-\exp(-cA_{\mathsf{log}})$, which is nearly 0 for large $c$, making $A_{d}^{T}\Delta_{d}\approx 0$ in Eq. 7.
We choose $c$ from $\{2,4,8\}$.
We then set $W_{\Delta}\approx 0$ and $b_{\Delta}=\mathsf{softplus}^{-1}(1)\approx 0.54$
in Eq. 4 so $\Delta_{d}\approx 1$.
This makes the state space layer close to its linear attention counterpart at initialization.

### 2. Correlated tokens should attend to each other

Having shown that state space layers can mimic
linear attention,
we now try to make them mimic attention layers that can copy,
such as the one in Fig. 2,
which implements a look-behind operation.
We focus on a single linear attention/state space layer,
*assuming* the layers before it learned a representation
amenable to copying.
Consider a copying example of length $n$,
where we have already copied $k<n$
of the $\mathtt{D}$-dim. tokens past the delimiter
$x_{\parallel}$
and want to copy the $(k+1)^{\mathsf{st}}$ one:
$X=(x_{1},\cdots,x_{n},x_{\parallel},x_{1},\cdots,x_{k})\in\mathbb{R}^{(n+k+1)
\times\mathtt{D}}$.
We assume that preceding layers $f$ have learned to superimpose
a position embedding as follows:

$$ $f(X)=(x_{1}+p_{1},\cdots,x_{n}+p_{n},x_{\parallel}+p_{1},x_{1}+p_{2},\cdots,x_ {k}+p_{k+1})=X+P\in\mathbb{R}^{(n+k+1)\times\mathtt{D}},$ $$

so that token with index $k$ in the paste string
will attend to token $k+1$ in the copy string
because $(x_{i+1}+p_{i+1})^{T}(x_{i}+p_{i+1})>0$,
assuming $x_{i+1}^{T}x_{i},x_{j}^{T}p_{j}\approx 0$ (uncorrelated)
and $p_{j}^{T}p_{j}=1$ (correlated).
That is, $f(X)^{T}f(X)\approx P^{T}P$ will
have similar structure to that in Fig. 2.
In this case, copying behavior will arise
in our state space/linear attention layer if
$P^{T}W_{C}^{T}W_{B}P\approx P^{T}P$,
i.e., when $W_{C}^{T}W_{B}\approx I$.
Since $W_{C},W_{B}$ are low rank ($\mathtt{N}<\mathtt{D}$), their product
cannot be exactly the identity;
using the fact that random Gaussian matrices are semi-orthogonal,
we could set $W_{C}\coloneqq W_{B}$ to get $W_{C}^{T}W_{B}\approx I$.
Initializing the queries and keys to be correlated
was also noted by Trockman & Kolter (), who suggest
these weights should not be strictly equal,
so we instead set $W_{C}\coloneqq\tfrac{1}{2}(W_{C}^{\prime}+W_{B})$.
In summary, assuming the model has learned a useful correlation
structure between tokens, setting $W_{C}^{T}W_{B}\approx I$
ensures this structure can be leveraged by attention.
For similar reasons, we experiment with initializing the convolution
in Mamba layers to the identity.

Figure: Figure 3: Testing the four components of our initialization on Mamba 1 & 2 for 10 seeds.
Refer to caption: x4.png

| Initialization | Purpose |
| --- | --- |
| $A\approx 1$ | Approximate<br>linear attn |
| $\Delta\approx 1$ |  |
| $W_{C}^{T}W_{B}\approx I$ | Encourage<br>recall |
| $\mathsf{Conv1d}\approx I$ |  |

### Which of these components matter?

In Fig. 3, we investigate the interaction of these
four possible mimetic initialization components, displaying all sixteen possible
off/on combinations. We investigate copying on 50-long strings and generalizing
to 100- and 300-long strings for a 24-layer Mamba with hidden size 1024
as in Jelassi et al. ().
For the $A$ and $\Delta$ initializations, we fix $c=8$ and $b_{\Delta}=0.54$.
For Mamba 1, we see that there is only a significant effect when
setting $A\approx 1$,
with no apparent benefit to setting $\Delta\approx 1$;
while setting $W_{C}^{T}W_{B}\approx 1$ has only a tiny effect,
using identity convolution initialization seems somewhat harmful.

For Mamba 2, we see a similar
advantage to using $A\approx 1$ initialization,
and a advantage to $W_{C}^{T}W_{B}\approx 1$
even without $A\approx 1$,
and the two interact to create even better models.
Adding identity convolution initialization
leads to much better performance still, reaching 100% accuracy
in many cases.
The positive interaction between $A\approx 1$
and $W_{C}^{T}W_{B}\approx 1$ and identity convolution
is especially apparent for 300-long strings.

The difference in the best initialization strategy
for the two architectures is likely explained by
the removal of linear blocks after the convolutional layer
in Mamba 2, as well as the addition of multiple state space heads.
Unless otherwise noted, we use the observations above
to determine our initialization strategy depending on the Mamba version:
For Mamba 1, we use $A,\Delta\approx 1,W_{C}^{T}W_{B}\approx I$,
and for Mamba 2 we add identity convolution initialization.

## 4 State Space Models want to be Transformers: Mimetic Initialization lets them get closer

Figure: Figure 4: Mimetic-initialized Mamba layers learn similar operations to Self-Attention layers in the same location *naturally* with no additional supervision on several tasks. Dotted lines: accuracy at training length (50), solid lines: generalizing to length 100.
Refer to caption: x7.png

Mimetic initialization leads to immediate and significant improvements
in copying ability.
In Fig. 1(b), we
can see that mimetic initialization allows a small 4-layer Mamba
to learn to copy strings with twice the training length with reasonable accuracy in just a few hundred steps,
which is far better than the tens of thousands of steps reported in previous work (Jelassi et al., ).
Note that mimetic initialization leads to Mamba learning
a state space layer whose attention map replicates the structure of
that of self-attention in Fig 2;
i.e., this layer has learned to (continue to) implement linear attention.
Mimetic initialization allows Mamba to quickly learn to copy from scratch.

### One mimetic init is all you need?

We continue our investigation of using mimetic initialization
to help Mamba learn recall tasks:
Given our observations that a single self-attention
layer is sufficient to learn these tasks to high fidelity,
and that a single Mamba layer can roughly approximate this
attention, we use mimetic init for *just one layer*
in the same position (Layer 4) of an 8-layer Mamba.

In addition to copying (Fig. 2),
we present results for additional three
synthetic tasks in
Fig. 4.
First, we investigate copying in stack order,
as
unpacking the compressed string in most-recently-added order
is potentially easier for SSMs.
Unlike normal copying, baseline Mamba is able to fit to the training length,
but it fails to generalize.
Mamba with mimetic init fits the training length much faster and generalizes better,
while the self-attention hybrid generalizes nearly immediately.
The story is similar for multi-query associative recall –
mimetic initialization leads to rapid learning and generalization to twice the length.
We also consider the sorting task, where tokens are sampled without replacement from a vocab of size 512.
Surprisingly, Mamba with mimetic init does even better than self-attention.
Mimetic initialization results in large improvements for all synthetic tasks considered.

Figure: Figure 5: Simple linear attention underperforms Mamba even for very high head dimension, especially at generalization. Dotted lines: accuracy at length 100, solid: at length 200; train length: 50.
Refer to caption: x10.png

### Is Mamba with mimetic init just linear attention?

In Figures 2 & 4,
notice that the mimetic initialized Mamba layer tends to mimic the
corresponding self-attention layer in the hybrid model;
the resemblance is clear for copying in normal and stack order.
For associative recall, it is less clear, but the Mamba layer
looks significantly more like it could implement a induction-head-like
function than typical Mamba layers.
Similarly, the interpretation is unclear for sorting, but the overall
structure matches.
At a high level, it seems like Mamba attempts to learn an approximation
to self-attention, but has much less capacity and sharpness.
Consequently, we ask if our initialization merely turns
state space layers into single-head linear attention layers.

In Figure 5, we present an ablation
study where we replace the target Mamba layer in our copying experiment
with simple causal linear attention with various head dimensions.
According to Eq. 12, we may expect mimetic init
to make Mamba layers equivalent to unkernelized linear attention layers
with head dimension equal to the state dimension.
Consequently, we compare Mamba with state size 32 to linear attention
with head dimension 32, which comes relatively close.
We plot generalization to $2\times$ and $4\times$-length in Fig. 5,
as the difference for fitting to the training length is small.
Nonetheless, Mamba still performs somewhat better than linear attention.
Linear attention performance depends on the head dimension,
with dimension 8 severely underperforming Mamba and dimension 1024
barely exceeding the performance of 32.
In contrast, doubling the state dimension of Mamba to 64
substantially improves generalization performance.
We visualize the difference in attention maps for the two operations;
we can see that Mamba’s is perhaps sharper/more consistent like that of self-attention.
Combined with better performance on copying, we conclude that mimetic init Mamba
layers are not *just* linear attention, but rather a related and superior
(for this task) non-linear operation.
The correlation between this “sharpness” and linear attention performance
has been exploited by recent work ().

## 5 Further Experiments on Mimetic Initialization

Mimetic initialization improves the recall abilities of Mamba 1 and 2 over a variety
of architecture settings and sequence lengths.
For all Mamba 1 experiments, we use state size 32, though we explore
different state sizes for Mamba 2, which has state size 128 unless otherwise noted.
For Mamba 2, we use head dimension 64 for all experiments.
All trials are for 5000 steps unless otherwise noted,
and we swept over a small set of learning rates;
our training pipeline is taken from Jelassi et al. ().
*Note:* While mimetic initialization has a strong effect size for Mamba 1, the architecture
generally struggles to copy for larger vocab sizes in the training lengths studied,
so we present Mamba 2 results for most larger-scale experiments in the paper.
Error bars are computed over five seeds.

Figure: Figure 6: Mamba 2 with mimetic init can learn to copy even for large vocabulary sizes.
Refer to caption: x11.png

### Vocabulary sizes

The larger the vocabulary, the more bits it should take to encode
content of a token to enable copying, and the harder it may be to memorize and copy the sequence.
While the previous work on copying focused on small vocabularies,
we showcase the ability of mimetic init to improve copying even for large vocabularies
in Fig. 6.
For Mamba 1, mimetic init allows decent copying performance up until a point,
and then degrades.
In contrast, baseline never learns to generalize.
For Mamba 2, mimetic init enables consistent $2\times$ length generalization
across sequence lengths, preventing the degradation with vocab size demonstrated by the baseline.

Figure: (a) State size vs. evaluation length
Refer to caption: x13.png

### State dimension

The copying ability of Mamba should be directly related to its state size,
according to Jelassi et al. ().
This allows Mamba to more easily approximate self-attention-like maps,
as we saw earlier.
We show this is indeed the case in Fig. 7(a).
Indeed, for baseline Mamba 2, perfect copying at training length 50 is only possible for sufficiently large state size.
However, if we use mimetic initialization,
the additional capacity from the state size is much more efficiently used,
and generalization (measured with the area under the curve) is far stronger –
$\mathtt{N}=32$ with mimetic init achieves performance comparable to
$\mathtt{N}=512$ with baseline init, a $16\times$ improvement in the use of capacity.
We show another view on this data in Fig. 7(b);
generalization length hardly grows with the log of the state size using baseline initialization,
while it grows linearly only after using mimetic initialization.
Mimetic init allows Mamba 2 to get closer to its true compression/copying capacity.

Figure: Figure 8: Mimetic initialization vs. Mamba 1/2 architecture sizes.
Refer to caption: x15.png

### Architecture size

In Figure 8, we investigate
mimetic init over different Mamba sizes (dimension, layers).
Surprisingly, a mere two layers seems to be sufficient,
with deeper networks improving generalization beyond $2\times$ length.
With embedding size 1024, Mamba 2 can copy very well for a variety of depths;
for multi-query associative recall, slightly deeper networks seem preferable.
In almost all cases, mimetic initialization leads to superior generalization
performance.

Figure: Figure 9: Mimetic init lets us nearly perfectly fit in-distribution even for long sequences on copying (left) and MQAR (right), and also boosts generalization performance (1024-dim 2-layer Mamba 2).
Refer to caption: x18.png

### Sequence length

Mimetic initialization lets us nearly perfectly fit to the training length
even for longer strings for both copying and multiquery associative recall (Fig. 9).
While baseline tends to struggle to learn to copy even
1000-long strings,
mimetic initialization allows fitting to around 4000-long
strings.
For MQAR, baseline breaks down around 900-long
strings, while mimetic initialization allows fitting to 1800-long or more.
The benefits apply for better generalization as well,
though Mamba still cannot strongly generalize to much longer strings than trained on.

Figure: Figure 10: Pretrained 768-dim. 24-layer Mamba 1 vs. from-scratch training (w/ mimetic init).
Refer to caption: x22.png

Figure: (a) The weights of some pretrained Mamba layers serve as a good initialization for copying, even when those weights are repeated uniformly for all layers in the “student” model. The layers that work well as an initialization for copying tend to have correlated $C$, $B$ weights and nearly-all-ones masks, such as Layer 31.
Refer to caption: x23.png

## 6 Comparing Mimetic Initialization to Pretraining

### Mimetic init mimics benefits of pretraining

We hypothesized that Mamba’s difficulty in copying may be an optimization issue
rather than fundamental capacity limitations.
That is, a Mamba that was first pretrained on
a general text corpus
may be a better representation of true copying abilities;
i.e., one should never train from scratch (Amos et al., ).
In Fig. 10, we see that finetuning a pretrained 130M
Mamba to copy or do associative recall on 50-character strings
results in good generalization, but
training from scratch with
mimetic init achieves similar results.
Note that the pretrained Mamba had a much longer $(>1\text{k})$ training length
than our from-scratch trials.
Considering this, our mimetic init results
get impressively close (esp. for shorter strings; dotted lines).

### Localizing the benefit of pretrained weights

Based on our linear attention observations,
the copying abilities of a pretrained Mamba may be localized to
a few layers, so we explore the capabilities of individual layers:
We use a pretrained teacher Mamba with layers $T_{i}:i\in[L]$, and then train $L$ student Mambas
where each of the $S_{j}:j\in[M]$ layers is initialized with $S_{j}\coloneq T_{i}$ for $i\in[L]$.
In this case, $L=48$ and $M=12$.
Using these pretrained weights can make it much easier to learn to copy (Fig. 11(a)),
but the effect size stands out for some particular layers, such as $T_{31}$.

We inspected the weights and attention maps of these layers
to see what might be behind the improved performance;
see some examples in Fig. 11(b).
Some layers such as $T_{31}$ look like our mimetic initialized layers,
with nearly all-ones average attention masks, correlated $W_{C},W_{B}$ weights,
and lower diagonal structure, similarly to self-attention layers in hybrid
Mambas earlier.
That is, the structure our initialization provides
seems to arise *naturally* in Mambas trained on sufficiently large and varied corpora,
and may be fundamental to Mamba’s copying and recall abilities.

## 7 Conclusion

We presented mimetic initialization for state space layers,
a simple and closed-form technique to greatly improve the copying and recall abilities
of state space models.
Mimetic initialization makes state space layers mimic linear attention
at initialization time,
and also mimics the structure of state space layers that contribute to copying and recall
abilities in pretrained models.
Our technique allows to estimate capabilities of SSMs more accurately, which have been alternatively over- and under-estimated in the literature (Jelassi et al., ; Waleffe et al., ).
Using a better initialization such as ours may assist
in developing new architectures starting from a smaller scale,
allowing for better predictions of their full-scale performance,
as is often done in practice in testbeds (Poli et al., ).
From a theoretical perspective,
our particular construction may provide insights into the tradeoffs between
state space layers and attention,
and may help to study the recall vs. non-recall capabilities of state space layers.
Improving the ability of state space layers to approximate attention
has already been noted in followup work to the original Mamba architecture (Dao & Gu, ),
and our initialization supports this concept.
More broadly and together with previous work on mimetic initialization,
our work helps to better understand pretraining,
to some extent disentangling its dual purposes of storing knowledge and serving as a good initialization.

## 8 Reproducibility Statement

We have provided all the necessary details
to reproduce our findings in the main text.
All experiments were done with multiple
random seeds, reporting the average and error bars.
We swept learning rates over
$\{0.001,0.0005,0.0001\}$.
We used the code from Jelassi et al. (),
found at [https://github.com/sjelassi/transformers_ssm_copy](https://github.com/sjelassi/transformers_ssm_copy),
and used pretrained weights
from [https://huggingface.co/state-spaces/mamba-130m](https://huggingface.co/state-spaces/mamba-130m)
and [https://huggingface.co/state-spaces/mamba-370m](https://huggingface.co/state-spaces/mamba-370m)
in some experiments.
For multiquery associative recall,
we used code from [https://github.com/HazyResearch/zoology](https://github.com/HazyResearch/zoology).
On two A100 GPUs, most training runs take
around 30-60m to complete,
with longer training times for very deep
models or those trained on very long sequences.
Source code will be released after publication.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Line 90, Why does valinna linear attention improve recall? This is not what has been observed in previous works and it is one of the main reasons for the relatively low adoption of vanilla linear Attention in modern architectures (linear Attention variants have been proposed since 2020, see [1, 2]). For example, the lack of recall capabilities of vanilla linear attention is one of the main motivations in [3] (this work is already mentioned in the manuscript). 
2. The paragraph starting in Line 354 is confusing (especially the sentence starting in Line 365). It is well know that Mamba layers are more expressive than Linear Attention as the latter is a special case of the former obtained when the state transition map does not fade (i.e. A=1), see for example [5, 6]. 
3. Line 202. What do the authors mean by “likely exploiting an implicit position embedding learned by the preceding Mamba layers”? Can they be more precise and point to an experiment showcasing that the Mamba layer learns positional encodings implicitly? 
4. Line 260. In past works [6, 7] it has been shown that such convolution is critical to further expand the state dimension of the SSM layer and that removing it could cause serious recall issues. Setting it to the identity (i.e. removing it) does not seem to be a good initialization in general. As this manuscript also mentions in Line 269. Do the authors have a theoretical motivation for the difference between Mamba 1 and Mamba 2 when it comes to the effects of the conv1d layer? 
5. The proposed method does not help for large scale experiments as the current initialization strategy is learned throughout training (line 520). Have the authors tested their initialization on larger scale language pre-training experiments? Do they expect their initialization to help when models are learned on larger scale experiments that are not only focused on copy/recall toy tasks?
6. Can the authors report the size of the state of the SSM models they studied in the paper? 


*Minor:*
Line 36 and Line 107, the Mamba layers are input dependent linear transformations, not time-dependent. 


[1] K. Choromanski et al., “Rethinking attention with performers”, arXiv preprint arXiv:2009.14794 (2020).

[2] S. Wang et al., “Linformer: Self-attention with linear complexity”, arXiv preprint arXiv:2006.04768 (2020).

[3] S. Arora et al., “Simple linear attention language models balance the recall-throughput tradeoff”, arXiv preprint arXiv:2402.18668 (2024). 

[4] A. Orvieto et al., “Resurrecting Recurrent Neural Networks for Long Sequences”, PMLR 202:26670-26698, 2023.

[5] S. Yang et al., “Gated Linear Attention Transformers with Hardware-Efficient Training”, arXiv preprint arXiv:2312.06635 (2023).

[6] L. Zancato et al., “B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory”. arXiv preprint arXiv:2407.06324 (2024).

[7] S. Yang et al., “Parallelizing linear transformers with the delta rule over sequence length.” arXiv preprint arXiv:2406.06484, 2024

[8] Y. Sun et al., “Retentive Network: A Successor to Transformer for Large Language Models”, arXiv preprint arXiv:2307.08621, 2023
soundness: 2
strengths: The paper addresses an important and timely question on how to improve recall capabilities of modern State Space Models (SSMs). The proposed initialization is thoroughly ablated with many toy experiments designed to test recall/copy abilities.
summary: The paper discusses the problem of improving recall/copying performance of State Space Models (Mamba). In particular, the paper claims that the difficulties in recalling from distant past instants of Mamba-like models are not due to the finiteness of their state and instead are due to poor initialization of their state transition maps. The main contribution of the paper is a novel structured initialization that allows state space layers to improve their recall capabilities.
weaknesses: While the proposed initialization scheme is sound, the novelty of the idea and technical contribution is incremental. Furthermore, the scope of the work is only limited to Mamba-like models and it is not generic to State Space Models (as stated in the title). To increase the scope of this work, given the relatively small scale experiments conducted, I’d suggest the authors extend their results to other SSM models like: GLA [5], LRU [4] and RetNet [8]. If only Mamba-like models are studied it would be good to modify the title since at the moment it reads as if it is applicable to “State Space Models” in general. 

The current manuscript is sometimes confused, SSM layers cannot mimic full Attention since the state of an SSM layer is fixed at the outset before seeing any sample, while the Attention’s state progressively increases as more and more samples are processed. This is a fundamental difference that cannot be bridged by keeping the same functional form of the SSM and only changing its initialization. Based on the results in the paper the proposed initialization avoids the collapse of the receptive field of the SSM’s state rather than mimicking Attention. 

Some important related references are missing, please see the Questions section below, can the authors discuss them? In particular, it is worth adding an explicit comparison with the analysis conducted in [4] which would help the reader contextualize the contribution of this work in the landscape of modern SSM models.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: * Did the authors try pretraining with this initialization scheme?
* Can the authors provide evidence that this initialization scheme would not lead to performance regressions on the basic language modeling tasks or non-recall based benchmarks?
soundness: 3
strengths: * The paper presents a wide array of experiments that vary state size, vocabulary size, sequence length, and embedding size to validate their results
* The paper is well written and easy to follow
* The paper provides a way to estimate the capabilities of SSM layers on copying tasks without expensive pretraining
summary: This paper provides an initialization scheme that allows SSM layers to learn copying tasks faster than the default initialization strategy.  They demonstrate that their mimetic initialization scheme allows SSM architectures to learn recall based tasks quickly.  They test their method on a variety of configurations, changing state size, embedding dimension, sequence length, and vocabulary size.
weaknesses: * This paper specifically trains for recall tasks, so it is unclear if this initialization scheme would lead to benefits for pretraining and not lead to performance regressions on non recall tasks
* Since hybrid models with attention layers that have inherent copying abilities combined with SSM layers are starting to become more prominent over pure SSM architectures, it is unclear if the copying abilities of the SSM components are of prime importance
* The paper does not provide any results for pretraining with this initialization scheme

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: * The plots showing attention patterns should be explained more clearly. What inputs were used to generate these visualizations? How dependent are they on the seed used for training? How do these patterns evolve over time? I think that this last question would be especially interesting given that the authors try to align SSM and self-attention at initialization, and I wonder whether this causes the dynamics to overall behave similarly, or rather that they simply start close and gradually diverge.
* I find the results in Figure 3 strange: first, it seems that the initialization has relatively small effect on Mamba 1(with accuracy going at best up to 50%) while the Figure 2 shows every different behavior (>90%). Am I misunderstanding something? Second, even the results for Mamba 2 seem to vary greatly based on seed, which is quite undesirable and detracts from the conclusions of the paper.
* Why do the authors believe there is such a large difference between Mamba 1 and Mamba 2? They mention this briefly on line 307, but I think this question deserves more attention.
* Similarly, why should Mamba perform better than Linear Attention, given that it was initialized to mimic its behavior?

Minor:
* Line 258: "who suggest these weights should not be strictly equal" — It may be useful to provide some motivation for this.
* Figure 8 mentions Mamba 1/2, but I believe only Mamba 2 is shown.
* l.485: I believe that the authors use a fixed T_i to initialize all S_j (since L and M are different) but this is not clear.
soundness: 2
strengths: The paper is well-written and addresses a fundamental limitation of SSMs compared to Transformers, which is still not fully understood. The proposed approach of initializing SSM layers to mimic self-attention is well-motivated, and the observed attention patterns suggest notable similarities in the behavior of the two. The authors present extensive experiments across various tasks and architectural configurations, confirming that the initialization generally improves performance. They also provide comparisons with linear attention and pretrained Mamba models.
summary: The paper investigates the challenges faced by State Space Models (SSMs), such as Mamba, on recall-based tasks compared to Transformer architectures. The authors propose an initialization technique that encourages SSM layers to behave more similarly to self-attention. They empirically show that this initialization significantly improves Mamba's ability to learn memory-related tasks.
weaknesses: While the authors motivate the problem well, I find that they do not spend enough time discussing the broader implications of their work. From a theoretical perspective, the initialization approach is mainly heuristic, and does not provide a fully satisfactory understanding of why SSMs struggle with such tasks. On the empirical side, it is unclear whether the proposed initialization could improve performance on language-based tasks. Additionally, I have some concerns about the robustness and clarity of the presented results (see questions below). The differing behaviors between Mamba and Mamba2 are also somewhat concerning, given that the motivation for the initialization does not distinguish between the two. To sum up, I would like the authors to clarify the main theoretical and empirical takeaways of their work.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Given the weaknesses listed above, I will recommend that this paper be rejected. While the proposed initializer is simple, intuitive, and seems effective, it lacks novelty, and the downstream analyses lack breadth. Furthermore, there is little theoretical justification in this paper (most of which appears to be drawn from prior work). This is acceptable only if a varied empirical analysis compensates for it, which is currently lacking. I think this paper would benefit significantly from another iteration. In particular, I would like to see the following addressed:

1. Explain why the line of work on orthogonal RNNs is not relevant here. If it is, please incorporate it into the paper. In particular, a more extensive related work section in the appendix would add substantial value.

2. Figure 2 is confusing. If the dots correspond to 10 random seeds, why not present the plot as an average with error bars? A more thorough multidimensional display could better reveal the meaningful variations across the ablated dimensions discussed later in the text.

3. The legend in Figure 5 is obscuring half of the plot.

4. Producing results that demonstrate utility beyond synthetic tasks would significantly improve the paper.

5. The authors note that Mamba outperforms linear attention on synthetic tasks but do not offer any intuition for this. Additionally, are the authors using the short convolution from the Mamba block in linear attention? If not, how is this a fair comparison? Please clarify

6. Much of the paper builds on linear attention, as this architecture is intended to serve as a prior in the mimetic initializer. If only one other SSM is to be benchmarked against, why choose Mamba? Isn’t GLA (or even S6) a more appropriate architecture for comparison? Mamba’s discretization step and short convolution are acting as unnecessary confounding factors. More generally, this initializer would be stronger if its effectiveness was demonstrated across a broader range of SSMs rather than just Mamba. Even so, could the authors at least please justify their choice of Mamba as the primary comparison model?

7. To what extent do the models deviate from the initialization prior after training?

8. In which regimes is the initializer most effective? Is it with many tokens to copy, long sequence lengths, etc.? Some plots potentially contain these results, but are hard to parse and there is no discussion in the paper.

9. Can the authors clarify the purpose of the MQAR analysis? Is it intended as a supplemental synthetic analysis, or does it show something distinct from copying? The paper’s claims primarily focus on copying, with MQAR briefly mentioned at the end.
soundness: 3
strengths: The main strengths of this paper lie in the simplicity of the proposed initializer and the robustness of the empirical results. In particular, I note the following:

1. The empirical findings are strong. Specifically, the mimetic Mamba models achieve significantly higher performance on both copying and recall tasks compared to the default Mamba models. They also reach this performance at a much faster rate.

2. The empirical evaluation is thorough. The authors conduct extensive sweeps across various vocabulary sizes, state sizes, and sequence lengths, demonstrating that mimetic Mamba consistently outperforms default Mamba across these settings.

3. The authors offer insight into how mimetic Mamba functions as a pretraining prior, providing intuition for why the initializer is effective.

4. The work appears reproducible based on the statement provided by the authors, although code has not yet been released.
summary: This paper addresses problems with copying (and more generally recall) in state-space models by proposing a novel "mimetic initialization" scheme based on the following two observations:

1. **State-space models as a generalization of linear attention**: The trainability of the $A_i$ matrices in state-space models allows them to act as a generalization of linear attention. However, this also can give rise to time decay that is too aggressive to learn copying tasks.
2. **Preserving correlation structure between input tokens**: By initializing the product of the $B$ and $C$ projection matrices to be the identity, the authors aim to maintain the correlation structure between input tokens.

By leveraging these two observations to construct an initialization scheme, the authors demonstrate that the resulting mimetic Mamba 1 and Mamba 2 architectures outperform the default Mamba 1 and Mamba 2 architectures on both copying and MQAR tasks. The authors extend the results across multiple task and model hyperparameters and conclude by demonstrating how mimetic initialization can be interpreted as an efficient alternative to model pretraining.
weaknesses: The weaknesses in this paper primarily lie in lack of presentation, novelty, breadth, and clarity. In particular, I note the following:

1. The paper is not particularly well-written. It seems rushed, and while the overall story makes sense, it lacks the logical flow both within and between sections that would make it easier for readers to follow. Some examples of this include the related works which would benefit from bolded paragraph headers like the authors use throughout the paper, rewriting the "Correlated tokens should attend to each other" section which is currently quite confusing to parse, moving the figures so they are actually near the text the author is reading, and more generally adding transitions between the sections that motivate why we need to do the next analysis so the story flows better. Also, I would consider moving the pretraining section after the first set of results since it provides nice intuition as to why the initializer works. 

2. The idea is not particularly novel compared to "Mimetic Initialization of Self-Attention Layers" (Trockman 2023), which is cited throughout the paper. Specifically, the authors adopt the concept of correlating the query and key projection matrices to maintain correlation structure amongst tokens from this prior work and simply apply it to the analogous components of Mamba.

3. The authors’ use of an identity prior on $A$ aligns with the intuition of orthogonal/unitary non-linear RNNs, a related line of work that is entirely ignored in the paper. Although the focus here is on linear systems, it seems incomplete not to mention this prior work. From my understanding, the approach presented in this work effectively reduces the rate at which the gradient vanishes over time, which is a longstanding idea in RNN research.

4. The authors restrict their results to synthetic tasks and do not demonstrate the applicability of this initializer (or any intuition derived from it) on language tasks, limiting the scope of the findings. Maybe something like the Pile is appropriate here?

5. The authors do not analyze other SSMs outside of Mamba (e.g., GLA, RWKV, Hyena, etc.), despite claiming that this initializer improves performance in SSMs more broadly. Intuitively, mimetic initialization should extend to something like GLA which has a distinct functional form from Mamba in the state-to-state dynamics. This would be nice to show. 

6. The authors benchmark this initializer only against the default Mamba initializer, even though the original Mamba paper presents other initializers that could have been included for comparison (these could be included in supplementary materials).

Meta-review: This paper studies the recall and copying performance of the MAMBA model to understand why it underperforms compared to transformer-based models. In turn, the authors propose a new initialization scheme aimed at improving recall performance.

The reviewers acknowledge that the paper addresses an important and not yet well-understood problem. The proposed initialization scheme is simple to implement and appears robust, and the paper is well organized and clearly written.

However, the experimental evaluation is limited, and one reviewer found some of the results puzzling. There are also concerns about the paper’s relevance given the rise of hybrid models in the field. (These concerns can be addressed, though.)

While the work is interesting, most reviewers feel it is not yet ready for publication. Although the authors provided additional results during the rebuttal phase, only minor revisions were made to the manuscript. In light of this, I agree with the reviewers that the paper is not ready for publication in its current form. I therefore recommend rejecting this submission, and suggest to resubmit a carefully revised manuscript to another venue.

Decision: Reject