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

Introduction

Transformers (Vaswani et al., 2017) are the workhorse of modern sequence modeling, achieving remarkable performance on a variety of tasks, but they have unavoidable inefficiencies. Specifically, they require $\Omega(L)$ memory111In some naive implementations of transformers, it is common to allocate a $L\times L$ matrix to compute the attention. However, memory efficient implementations, such as FlashAttention (Dao et al., 2022), compute the attention with $O(L)$ memory. and compute to predict the next token of a sequence of length $L$.

This has spurred a boom in attempts to create architectures that can achieve similar performance as transformers, but with $O(1)$ memory to predict each token. This class of models includes state space models like S4 (Gu et al., 2021) or Mamba (Gu & Dao, 2023), as well as traditional RNN models (Hochreiter & Schmidhuber, 1997) and models that can be trained in parallel like linear attention (Katharopoulos et al., 2020; Choromanski et al., 2020) and parallel RNNs (Bradbury et al., 2016; Peng et al., 2023; Sun et al., 2023). In this paper, we will refer to this entire class of models that use a fixed-size memory as “generalized state space models” or GSSMs (see a formal definition in Section 2).

Recent work has demonstrated impressive performance of GSSMs, but it is not yet clear what these models sacrifice for their improved efficiency, if anything. In this paper, we find that one particular capability that is sacrificed is the ability to retrieve and repeat parts of the input context. As a result, transformers are better than GSSMs at a variety of tasks that require accessing arbitrary parts of the context.

To understand this gap in capabilities, we begin by presenting a theoretical analysis of the copying task222Note that we study copying of the input and not copying of training data (McCoy et al., 2023; Carlini et al., 2022). First, we show via construction that a simple transformer model can copy strings of length that is exponential in the number of heads of the transformer. This construction relies on the ability of the transformer to implement a mechanism of “storage” and retrieval of sequences of n tokens (n-grams), where the n-grams are used to track where to copy from. In contrast, we show that, trivially, GSSMs cannot accurately copy strings with more bits than the size of the latent state.

Pythia: 410M 1.4B 2.8B Mamba: 360M 1.4B 2.8B

Our theory studies representation expressivity, but not whether these representations will be learned. Moreover, in practice a large GSSM may have enough capacity to represent the entire input in the latent state, at least in theory. To resolve these concerns, we conduct a variety of synthetic experiments with models of $\sim$160M parameters. We find that transformers are both much more efficient at learning to copy (Figure 1(a)) and also generalize better to longer inputs (Figure 1(b)). Additionally, we verify experimentally that the copy “algorithm” learned by transformers indeed relies on n-grams to perform a lookup of where to copy from (Figure 3), similarly to our theoretical construction.

Finally, we present a variety of experiments on pre-trained models to test their ability to remember and access the input context. In particular, we show that Pythia transformers (Biderman et al., 2023) outperform Mamba GSSMs (Gu & Dao, 2023) of similar size at a variety of memory-intensive tasks including copying and retrieving information from the context (Figure 1(c)). This is especially notable since the Mamba models achieve lower perplexity than the Pythia models at language modeling on the Pile (Gao et al., 2020). These experiments illustrate the practical relevance of the memory issues that we raise, and hint at one way that architectual choices can impact the downstream performance of LLMs above and beyond training perplexity.

Theory: Representational Capacity

In this section we use the copy task for a theoretical comparison between state space models and transformers. We prove two main results. First, we construct a small transformer that solves the copy task for sequences lengths that are exponential in the transformer size. Second, we show that any state space model fails to solve the copy task, unless its latent state grows linearly with the sequence length.

Transformers.

The copy task.

where $H_{1:L}(\cdot)$ denotes the first $L$ tokens generated by $H$. That is, we expect the model to output an exact copy of ${\bm{x}}$.

2 Transformers can copy inputs of exponential length

In this section we show that transformers can implement the copy operation for input sequences with length exponential in the number of heads. Namely, we construct a transformer with two blocks that gets small error on the copy task.

The key idea in the construction is to first “hash” sequences of $n$ tokens ($n$-grams), then at each iteration of the auto-regression attend to the previous occurrence of the most recent $n$-gram, and output the succeeding token. That is, we show that a transformer can implement the copying algorithm illustrated in Figure 3 (and see also Algorithm 1 in the Appendix).

Positional embedding: Hard-ALiBi.

Guarantees.

The copy algorithm given in Algorithm 1 (and similarly, our transformer construction) can perfectly copy the input sequence, as long as there are no repeated $n$-gram patterns in the input. Therefore, the error of the algorithm depends on the probability of repeated $n$-grams:

Below we state the main theoretical result on copying with transformers, showing that transformers can copy their input, with error bounded by the probability of repeated $n$-grams:

Intuitively, the probability of repeated $n$-grams decays quickly when increasing the value of $n$. Indeed, we show that for the uniform distribution over sequences, this probability decays exponentially with $n$:

Combining the above results, we get that transformers can copy sequences of tokens drawn from the uniform distribution, using a number of parameters that depends only logarithmically on the input sequence length.

For simplicity we do not limit the precision of the parameters or activations, but note that our results hold for finite-precision transormers, using $O(\log(\log(L)))$ bits.

3 State Space Models cannot copy inputs beyond memory size

We saw that transformers are able to copy uniform sequences of tokens, with parameter count logarithmic in the sequence length. We now show that GSSMs cannot copy uniform input sequences, unless the capacity of their state space grows linearly with the size of the sequence length. This is intuitive: to be able to copy the entire input sequence, the model needs to store it in its state space, which requires the memory to grow linearly with the sequence length.

Given Theorem 2.7, the following Corollary is immediate:

Learning to Copy

In the previous section, we proved that transformers can represent the copy operation for exponentially long sequences, while GSSMs fail to copy long sequences due to their limited memory. While these results show that in theory, transformers can outperform GSSMs, our theoretical results do not establish that such a gap will be observed in practice for two reasons. First, it is not clear that transformers can indeed learn to copy from examples. Second, GSSMs in practice may use a large latent state memory, so that our bounds only hold for very long sequences of tokens. For example, a latent state of 1000 32-bit floating point numbers has enough bits to store at least 2000 tokens from a 50K token vocabulary. However, even though a GSSM could fit the context into memory, it may not learn to do so.

Our goal in this section is to verify that our theoretical analysis bears out experimentally when training models from scratch on synthetic data, before moving on to study pretrained models in the next section. Specifically, we train transformers and GSSMs (LSTM (Hochreiter & Schmidhuber, 1997) and Mamba (Gu & Dao, 2023)) on variants of the copy task shown in Figure 2.

We now provide a brief overview of our experimental setup. Further details may be found in Appendix A.

In all our experiments, we set the model hyperparameters so that the Mamba and transformers have a similar number of parameters ($\approx 160$ million parameters). Since we find that large LSTMs are hard to train (as confirmed in Pascanu et al. (2013)), we use the largest LSTM we managed to train which has $\approx 40$ million parameters.

Dataset.

During training, we generate in an online manner a batch of 64 examples at each epoch. At test time, we evaluate our models on $10$ batches of $128$ examples. We report the mean and standard-deviation over these 10 batches. If not specified otherwise, our token space $\mathcal{V}$ is of size 30 and made of the alphabet letters i.e. $\mathcal{V}=\{a,\dots,z,\left\langle{\tiny\textsc{BOS}}\right\rangle,\left\langle{\tiny\textsc{EOS}}\right\rangle,\left\langle{\tiny\textsc{COPY}}\right\rangle\}$ where $\left\langle{\tiny\textsc{BOS}}\right\rangle$ is the beginning of sentence token, $\left\langle{\tiny\textsc{EOS}}\right\rangle$ the end of sentence token and $\left\langle{\tiny\textsc{COPY}}\right\rangle$ the separator token. All the strings are sampled uniformly i.e. we first sample the length of the sequence and then independently sample each position of the string from $\mathcal{V}$. Finally, we “pack the context” with i.i.d. sequences during training similarly to (Zhou et al., 2023): we fill the context with multiple independent samples of the task.

Positional information.

Positional information also plays an important role in the length generalization capacity of Transformers (Jelassi et al., 2023; Kazemnejad et al., 2023; Shen et al., 2023). Previously popular methods of input-layer positional embeddings (e.g. sinusoidal (Vaswani et al., 2017) or learned (Radford et al., 2019)) have been replaced by relative positional encodings at each attention layer (e.g. RoPE (Su et al., 2023), Alibi (Press et al., 2021), or NoPE (Kazemnejad et al., 2023)). Below, we experiment these positional encodings along with the Hard-Alibi encoding introduced in Section 2.

2 Data efficiency on the copy task

We begin by training our models on the simple task of copying a sequence of input tokens described in Figure 2. The model gets an input of $\leq L$ tokens followed by a Separator ($\left\langle{\tiny\textsc{COPY}}\right\rangle$) token, and needs to output the same sequence again from the beginning. In this section, we focus on in-distribution learning: we train on strings of random length $\leq L=300$ and record the string-level accuracy on evaluation strings sampled from the training distribution.

Results for this experiment are shown in 1(a). Clearly, there is a large gap between the transformers and GSSMs. We observe that the transformers need 100x less samples than the best GSSMs to learn the copy task.

Note that the sharp changes in accuracy displayed in 1(a) are due to the log-scaled x-axis and choice of string-level accuracy as a metric. In 9(a), we report the character-level accuracy, which yields smoother curves demonstrating the learning process of GSSMs. Regarding LSTMs, we find that they do not manage to learn on length-300 strings even at the character level. In 9(b), we show that LSTMs are able to learn to copy on shorter strings and that string length is the bottleneck.

3 Length generalization on the copy task

The prior experiment demonstrates superior efficiency of learning in-distribution. Now, we test the ability of the learned functions to generalize out-of-distribution. Specifically, we consider generalization from short sequences to longer sequences. Testing this sort of generalization can help us to better understand which function the model has learned, i.e. whether the model has truly learned the “correct” copy operation or whether it just learned to copy sequences of the particular size it was trained on.

Here, we train all models on sequences of $\leq 50$ tokens, and test them on sequences of up to $1000$ tokens, reporting string-level accuracy. As seen in 1(b), all models are able to (eventually) solve the task in-distribution on lengths of $\leq 50$, but transformer-based models display much better generalization to longer inputs compared to GSSMs. Namely, we observe that the performance of the GSSMs (LSTM and MAMBA) drops to zero almost immediately when increasing the input length, while the performance of transformers decays much more gradually with length.

When looking at the relative performance of different transformer models in 1(b), it becomes clear that the positional encoding is important to length generalization. Specifically, the ALiBi and NoPE transformers dramatically outperform the RoPE model on longer inputs. This is likely because the sinusoidal embeddings of RoPE create a more dramatic change than the decay of ALiBi or NoPE when we go to longer inputs.

Improved generalization with Hard-ALiBi.

To test our understanding of how transformers learn to copy, we now consider swapping in the Hard-ALiBi positional encoding that we used in our theoretical construction of hash-based copying (introduces in Subsection 2.2 and illustrated in Figure 8 in the Appendix). 1(b) shows that a transformer trained with Hard-ALiBi embedding on sequences of length $\leq 50$ achieves almost perfect length generalization up to sequences of length 1000. Note that this is well beyond the context length ever encountered in training.

4 Transformers learn to use n-gram hashing

Next, we attempt to determine whether the transformer trained on the copy task indeed applies the mechanism of storage and retrieval of n-grams. To do this, we evaluate the performance of a transformer with Hard-ALiBi positional encoding trained on the copy task when tested on a distribution of examples that intentionally contains duplicate n-grams. That is, we draw uniform sequences of tokens, and then randomly replace some n-gram with another n-gram that already appears in the sequence, such that each example always contains two copies of the same n-gram (typically followed by a different token). We use the Hard-Alibi model here since it performs the best for the copy task as showed in 1(a). Figure 4 shows the performance of the transformer for different choices of $n$. We observe that the transformer maintains roughly the same accuracy for $n\leq 4$, but that its accuracy starts dropping when the inputs contains duplicate sequences of 5 or more tokens. This suggests that the transformer relies on something like 5-gram retrieval to do the copy task.

5 GSSMs cannot arbitrarily retrieve from context

Transformer: NoPE Alibi HAlibi GSSM: LSTM Mamba

Transformer: NoPE Alibi HAlibi GSSM: LSTM Mamba

We now introduce another task to probe the mechanisms that the models use to copy from the context: the n-gram lookup task. In this task the model needs to use a given n-gram as a key to look up the k-token key that follows the query. We consider two variants of the task: suffix keys and prefix keys. In both variants, we assess length generalization to understand the function that the models have learned.

First, we consider the suffix key version of n-gram lookup. In this task, the model is given a sequence $L$ of input tokens, a separator, and then an n-gram from the input sequence. The model then needs to output a sequence of $k$ tokens following the chosen n-gram (see Figure 5 for an illustration). This task is closely related to induction heads (Olsson et al., 2022). This task requires the model to be able to “store” the entire context in order to effectively find the correct key to access it’s query. We train all models on sequences of at most 30 tokens and show results in Figure 5. Transformers perform well on this task, with a relatively small drop in performance when increasing the sequence length up to 100. This suggests that transformers can learn to perform n-gram storage and retrieval. GSSMs, however, perform poorly beyond their training distribution. Intuitively, this task still requires the models to store the entire input sequence, something that GSSMs struggle to do.

Next, we try the prefix key version of n-gram lookup. Here we provide the n-gram key at the beginning and then the full input sequence (illustrated in Figure 6). In this version of the task the model does not need to store the entire input since it can look for the key on the fly as the sequence is processed. This is good for the GSSMs, since they can write the key into the state and then ignore inputs that do not match. Indeed, GSSMs achieve perfect length-generalization on this variant. Interestingly, the GSSMs even outperform the NoPE and ALiBi transformers (although not the Hard-Alibi model). We hypothesize that this may be an issue where these positional embeddings make it more difficult to effectively perform the hashing lookup over a long distance in relative positions. Taken together, these results illustrate how GSSMs seem to be memory limited, but can be effective when the tasks only require a summary of the inputs rather than storing the entire context.

Pre-trained Models

In this section, we compare the performance of pre-trained transformers and pre-trained GSSMs on memory-intensive tasks such as copying long strings, retrieval and few-shot question answering. We show that transformers outperform GSSMs of similar scale on such memory-intensive tasks, even when the GSSM has lower perplexity as a language model. These results confirm that the limitation of GSSMs raised in previous sections apply to large scale models trained on real pretraining data.

In the experiments below, we compare Pythia transformer models (Biderman et al., 2023) of sizes ranging from 410M to 2.8B against Mamba models (Gu & Dao, 2023) of similar sizes. All these models have been pre-trained on the Pile (Gao et al., 2020) and use the same tokenizer. The Mamba models generally have slightly lower perplexity on the training set for a given size. The main difference between the Pythia and the Mamba models is their architectural design.

We compare these models by measuring their performance while varying the input instance length and consider two types of tasks: copy-based and information retrieval tasks. The copy-based tasks consist of presenting a random text to the model and asking it to copy the text. In the information retrieval tasks, we provide a text to the model and ask it a related question. These retrieval tasks can be seen as ”selective copy”, since the model needs to copy a small chunk of the input text in order to respond to the question. To measure performance, we use the string-level accuracy in all the experiments except in 7(c) where we consider question answering and thus report the F1 score. We evaluate the models over 10 batches of size 64 for all the tasks except for question answering where we evaluate over 50 questions because the number of questions with a given context length is limited. Further details are in Appendix A.

Pythia: 410M 1.4B 2.8B Mamba: 360M 1.4B 2.8B

Pythia: 410M 1.4B 2.8B Mamba: 360M 1.4B 2.8B

2 Copying the input context

We first observe that pre-trained transformers outperform pre-trained GSSMs at copying long natural language strings. In 7(a), we randomly sample strings from the C4 dataset (Raffel et al., 2020) with varying number of tokens. Our prompt consists of two copies of the sampled string plus the first word of the string and we expect the model to complete the third copy. Even the smallest transformer model dramatically outperforms the largest GSSM. This happens even though the large GSSMs have enough bits in the state variable to potentially store the context. This confirms the idea that this is an architectual bias of transformers that makes it easier for them to copy from the context.

Unlike strings of tokens sampled uniformly at random, natural text can often be compressed, possibly allowing language models to copy longer strings even with limited memory. To test whether this matters, in 7(b) we conduct the same experiment as above but randomly shuffle the order of the words in the strings. We find that when we shuffle the words, both GSSMs and transformers perform worse on the task, but the effect is more stark for GSSMs. Even the largest GSSM now gets zero accuracy on strings of length 300. This suggests that when the input is more difficult to compress, the GSSM suffers due to its fixed size state.

3 Retrieval from the input context

While copying provides a clear task to separate the model classes, it is not a particularly realistic task. That said, it presents an extreme case of a type of behavior that is highly relevant for many tasks of interest. In particular, many tasks require retrieving specific information from the context that is relevant to the desired output. This subsection presents examples of how our results transfer to more practical tasks.

We first consider a “phone-book” experiment where we provide a synthetic phone-book to the model and ask it to return the phone number when given a name. We generate the phone-book by randomly sampling $L$ names and their associated phone number. One line of this phone-book looks like “John Powell: 609-323-7777”. Our prompt to the model consists of the phone-book, two few-shot examples and a question asking for the phone number of a randomly sampled name from the phone-book. 1(c) reports the accuracy obtained by the pretrained transformers and GSSMs while varying the size of the phone-book $L.$ We observe that even the smallest transformer (410M parameters) outperforms the largest GSSMs (2.8B parameters) when the phone-book size is long enough ($L\geq 70$). This shows that in retrieval tasks which require access to the whole context, GSSMs struggle to store the relevant information in their fixed-size state.

Question-Answering.

In this experiment, we compare the 2.8B parameter Mamba and transformer models444In our experiments, smaller models were unable to achieve reasonable and consistent performance on this dataset., on the SQuAD question-answering dataset (Rajpurkar et al., 2018). This dataset provides text paragraphs together with a few questions regarding the text. We probe the models to answer the question by providing a single demonstration of a question/answer pair (corresponding to the same text) before giving the target question. We bin the paragraphs according to their lengths, and report the F1 score as a function of the paragraph length for both models in 7(c). We observe that while for short paragraphs, both the Pythia transformer and Mamba achieve comparable performance, the performance of Mamba degrades more quickly with the paragraph length, while the transformer-based model maintains a similar accuracy even for longer texts. This result shows that the fixed-memory of GSSMs also limits their performance on standard natural tasks.

Related Work

There exists a broad body of prior work on the representational capacity of GSSMs like RNNs (Merrill, 2019; Merrill et al., 2020) as well as transformers (Weiss et al., 2021; Merrill et al., 2022; Wei et al., 2022; Sanford et al., 2023; Edelman et al., 2022). Previous works that study transformers do so through comparison to other complexity classes, such as threshold circuits (Merrill et al., 2022), RASP language (Weiss et al., 2021) or first-order logic (Chiang et al., 2023) (see Strobl et al. (2023) for a thorough review). These works do not provide insights into how transformers implement algorithms for solving specific problems. In contrast, our theoretical result constructs a transformer for the copy task, which illustrates the mechanism and provides tight bounds on the model size. Together with the result showing that GSSMs cannot copy long sequences, our theory characterizes the power of different sequence models on the copy task. Other theoretical separation results between transformers and RNNs (Sanford et al., 2023; Merrill, 2019) use more complex tasks of less practical relevance.

Other papers have previously demonstrated the capacity of transformers to leverage the entire input context for tasks like retrieval, question answering, and in-context learning (Devlin et al., 2018; Raffel et al., 2020; Petroni et al., 2020; Brown et al., 2020; Liu et al., 2023b). Another line of work has studied the “induction head” mechanism in transformers that performs a retrieval operation much like the one we observe for copying (Olsson et al., 2022). But, to our knowledge, there is not a comparison in related work between transformers and GSSMs of similar quality on these tasks.

Several of our experiments study length generalization as a way to assess whether the model found the “right way” to solve the task. Prior work on length generalization in transformers has focused on the data distribution (Anil et al., 2022), positional embeddings (Kazemnejad et al., 2023), and arithmetic tasks (Delétang et al., 2022; Ruoss et al., 2023; Jelassi et al., 2023; Zhou et al., 2023). We extend many of these ideas to the copying task.

Finally, we note that while we focus on tasks where transformers outperform GSSMs, there are also tasks where GSSMs outperform transformers. For example, Liu et al. (2023a) shows that transformers fail to generalize out of distribution for “flip-flop language modeling”, while LSTMs do so easily. These tasks require tracking a small $O(1)$ state variable over time. Another benefit of GSSMs is the ability to input long contexts like DNA sequences that may be impractical for transformers (Nguyen et al., 2023).

Discussion

We have demonstrated through theory and experiments that transformers are better than GSSMs at copying from their input context. However, we emphasize that state space models have many advantages over transformers. The memory and computational complexity of GSSMs does not increase with the input length, which is ideal for training and inference on long inputs. Additionally, state space models such as RNNs are better at tracking state variables across long sequences (Liu et al., 2023a), which may be useful for generating long consistent text. Importantly, language processing in the human brain appears to be much more similar to how state space models process language (Tikochinski et al., 2024). We therefore believe that future work should focus on building hybrid architectures that endow state space models with an attention-like mechanism, allowing them to retrieve relevant pieces of text from their input. Indeed, humans have an incredibly limited capacity for memorizing sequences (Miller, 1956), but can translate entire novels if we allow them to look back at the text (Shelton, 1612).

Acknowledgements

We thank Boaz Barak for helpful discussions. Kempner Institute computing resources enabled this work. Samy Jelassi acknowledges funding supported by the Center of Mathematical Sciences and Applications. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Sham Kakade acknowledges funding from the Office of Naval Research under award N00014-22-1-2377.

References

Appendix A Experimental setup

In this section, we provide additional details about our experimental setup. We first give a description of the positional encodings used in our transformers experiments (Subsection A.1) and then give details about the training and evaluation procedures (Subsection A.2).

We consider multiple positional encoding schemes in our experiments in Section 3:

the NoPE scheme (Kazemnejad et al., 2023) where no positional information is added to any of the attention scores (8(a)). This architecture choice helps to get better length generalization in multiple tasks including the copy task.

the ALiBi scheme (Press et al., 2021) which biases the attention scores with a penalty that is proportional to their distance (8(b)). $m$ is a head-specific slope fixed before training.

the Hard-ALiBi scheme introduced in Section 2 which has $M$ masked attention heads where we explicitly force the model to attend to their directly previous tokens and $H-M$ heads set to be NoPE attention heads. In 8(c), we display the case where we have $M=4$ masked heads: in the first head, the tokens just attend to themselves; in the second head, the tokens attend to themselves and to previous ones; in the third head, the tokens attend to themselves, the previous ones and the second preceding tokens. The remaining $H-M$ heads are set to NoPE.

A.2 Pretraining and evaluation details

We implement all of our training in Pytorch (Paszke et al., 2019). We use the HuggingFace library (Wolf et al., 2019) and the Mamba GitHub repository (Gu & Dao, 2023).

Architectures.

In our experiments in Section 3, the backbone of our transformers is the GPT-NeoX architecture. We set the number of layers to 12, the hidden size to 1024 and the number of heads $H=16$. We consider the different positional encodings that are described in Subsection A.1. For Alibi, we set the head-specific scalar as in the original paper i.e. $m_{h}=2^{-h/2}$ for $h\in\{1,\dots,H\}.$ For the Hard-Alibi model, we sweep over the number of masked heads $M\in\{2,\dots,10\}$ and found that the best model corresponds to $M=6$. Regarding the Mamba models, we set the number of layers to 24 and the hidden size 1024. We also sweep over the state space dimension $S\in\{16,32,64,128,256\}$ and found the best model is $S=32$. This choice of hyperparameters ensures that both transformers and Mamba models have a comparable number of parameters. Lastly, our LSTM is made of 4 layers and width 1024.

Training hyperparameters.

In Section 3, at each epoch, we sample online a batch size of size 64. We fill the context with examples so we choose a context length ($C=420$ for all the experiments except 1(a) where we set $C=620$) and pack as many examples as possible to fit this context. So in our case, one sample contains many instances. We run the experiments for 15 epochs for both transformers and Mamba while for LSTMs we need 300 epochs. All methods are trained with the AdamW optimizer (Loshchilov & Hutter, 2017) with learning rate 5e-5, a linear rate decay schedule, 300 steps of warmup and default weight decay of 1e-1. Finally, to train all the models, we use the next-token prediction loss but we apply a mask on the input instance so that we only penalize the model whenever it makes a mistake on the labels (and not on the inputs and labels jointly).

Compute resources.

Pretraining was all done on an internal cluster using RTX8000 GPUs. We estimate that the final training run needed to produce the results in the paper took approximately 600 GPU hours.

Evaluation algorithm.

We evaluate the models over 10 batches of size 64 for all the tasks except for the question answering one where we evaluate over 50 questions because the number of questions with a given context length is limited.

Decoding algorithm.

At inference, all our models use greedy decoding for generation and we set the temperature to 0.

Appendix B Additional Experiments

In Subsection B.1, we focus on the in-distribution learning of the copy task and show that the number of samples needed by GSSMs is much higher than the one for transformers. In Subsection B.2, we study the performance of pre-trained models on the copy task in the case where the strings are sampled uniformly. This experiment shows that when the text to copy is totally random, the gap between pre-trained transformers and GSSMs is even larger.

Transformer: RoPE NoPE Alibi HAlibi GSSM: LSTM Mamba

In this section, we provide additional plots to complement the data efficiency experiment from 1(a). We want to highlight the following points:

in 1(a), we see a sharp transition for the Mamba learning curve. However, 9(a) shows that the learning process is more smooth at the character level. Besides, LSTMs are not able to learn the copy on length-300 strings even at the character level.

We consider the experiment of learning to copy much shorter strings namely strings with length $\leq 30$. 9(b) shows that the gap in terms of training examples between transformers and Mamba is much smaller i.e. Mamba only needs 10x more data. Besides, we see that the LSTM is able to learn the copy task but it needs 100x more data than transformers.

B.2 Pre-trained models on the uniform copy task

In this section, we provide an additional experiment that shows the superiority of pre-trained Pythia over pre-trained Mamba models in the copy task.

Pythia: 410M 1.4B 2.8B Mamba: 360M 1.4B 2.8B

We consider the same setup as in Section 3: we sample uniform strings of alphabet characters with a fixed length and ask the model to copy it by using the same prompt format as the one described in Subsection 4.2.

This setting is a more extreme version of 7(b) since the strings are more random: in 7(b), the order of the nouns were random but the nouns were English nouns while in 7(b), the strings are totally random. In Figure 10, we see a clear separation between the transformers and Mamba models with the smallest Pythia outperforming the largest Mamba. However, compared to 7(b), the Pythia performance is much higher since the 1.4B model able to get almost 100% accuracy.

Appendix C Proofs - Upper Bound

This section gives a detailed proof of Theorem 2.3 and Lemma 2.4.

We begin by introducing some technical lemmas that we use in the proof of Theorem 2.3.

𝑛2…𝐿i=n+2,\dots,L do $k_{i}\leftarrow s(x_{i-n},x_{i-n+1},\dots,x_{i-1})$ $v_{i}\leftarrow x_{i}$ end for for $j=1,\dots,L$ do if $j\leq n$ then $y_{j}\leftarrow x_{j}$ else $q_{j}\leftarrow s(y_{j-n},\dots,y_{j-1})$ Let $i\in[L]$ s.t. $k_{i}=q_{j}$, and set $y_{j}\leftarrow x_{i}$ end if end for Output: sequence $y_{1},\dots,y_{L}$ Lemma C.1. Let $h_{t}({\bm{x}}_{1},\dots,{\bm{x}}_{i})=\frac{1}{\min(t,i)}\sum_{j=\max(1,i-t+1)}^{i}{\bm{x}}_{j}$. Then, $h_{t}$ can be computed using a hard-ALiBi attention head.

Let $W_{k},W_{q}=0$ (zero matrix) and let $W_{v}=I_{d}$ (indentity matrix). We choose $b_{i}\in\{0,-\infty\}^{i}$ s.t.

Assume that $d=\left\lceil\log(D)\right\rceil+2$. Then, there exists an embedding $\Psi$ s.t.

For $x^{\prime}\neq x$ it holds that $\left\langle x,x^{\prime}\right\rangle<1-\frac{1}{d}$.

For every $x\neq\left\langle{\tiny\textsc{BOS}}\right\rangle$, $\left\langle\Psi(x),\Psi(\left\langle{\tiny\textsc{BOS}}\right\rangle)\right\rangle=0$, and for every $x\neq\left\langle{\tiny\textsc{COPY}}\right\rangle$, $\left\langle\Psi(x),\Psi(\left\langle{\tiny\textsc{COPY}}\right\rangle)\right\rangle=0$.

Denote $d^{\prime}=\left\lceil\log(D)\right\rceil$, and observe that we can encode all $D$ “non-special” tokens as vectors in $\left\{\pm\frac{1}{\sqrt{d}}\right\}^{d^{\prime}}$, and denote this encoding by $\Psi^{\prime}$. Now, define:

1𝐾𝑏𝑎11𝐾𝑏𝑎\displaystyle s_{i}=\frac{\exp(z_{i})}{\sum_{j=1}^{K}\exp(z_{j})}\geq\frac{\exp(a)}{\exp(a)(1+K\exp(b-a))}=\frac{1}{1+K\exp(b-a)} Finally, for every $j\neq i$:

C.2 Proof of Theorem 2.3

Claim: $\hat{g}_{t}(\Psi(x_{1}),\dots,\Psi(x_{i}))=\bm{1}\{i>n\}\cdot\Psi(x_{i-t})$

Proof: Fix some $j\in[d]$. Observe that for all $i$, $\left\lvert e_{j}\cdot g_{t}(\Psi(x_{1}),\dots,\Psi(x_{i}))\right\rvert\leq 1$.

If $i\leq n$, we have $g_{n}(\Psi(x_{1}),\dots,\Psi(x_{i}))=\frac{1}{i}\sum_{j^{\prime}=1}^{i}\Psi(x_{j^{\prime}})$ and so $\Psi(\left\langle{\tiny\textsc{BOS}}\right\rangle)\cdot g_{n}(\Psi(x_{1}),\dots,\Psi(x_{i}))=1$ where we use the properties of $\Psi$ and the fact that $x_{1}=\left\langle{\tiny\textsc{BOS}}\right\rangle$. Therefore, $\hat{g}_{t,j}(\Psi(x_{1}),\dots,\Psi(x_{i}))=0$.

where we use the fact that $x_{i-t}\neq\left\langle{\tiny\textsc{BOS}}\right\rangle$ and therefore $\Psi(\left\langle{\tiny\textsc{BOS}}\right\rangle)\cdot\Psi(x_{i-t})=0$.

Proof: Denote $g_{t,i}=g_{t}(\Psi(x_{1}),\dots,\Psi(x_{i}))$ and $h_{1,i}=h_{1}(\Psi(x_{1}),\dots,\Psi(x_{i}))$. Observe:

For all $t=1,\dots,n$, define $g^{*}_{t}({\bm{x}}_{1},\dots,{\bm{x}}_{i})=\sigma(\Psi(\left\langle{\tiny\textsc{COPY}}\right\rangle)\cdot g_{t-1}({\bm{x}}_{1},\dots,{\bm{x}}_{i}))$.

Claim: $g^{*}_{t}(\Psi(x_{1}),\dots,\Psi(x_{i}))=\bm{1}\{i=L+t\}$

Proof: Denote $g_{t,i}=g_{t}(\Psi(x_{1}),\dots,\Psi(x_{i}))$. Observe:

If $i=L+t$ then $g_{t-1,i}=\Psi(x_{i-t+1})=\Psi(x_{L+1})=\Psi(\left\langle{\tiny\textsc{COPY}}\right\rangle)$ and therefore $g^{*}_{t,i}=1$.

If $i\neq L+t$ and $i>t-1$ then $g_{t-1,i}=\Psi(x_{i-t+1})\perp\Psi(\left\langle{\tiny\textsc{COPY}}\right\rangle)$ and therefore $g_{t,i}^{*}=0$.

If $i\leq t$ then since $x_{1},\dots,x_{i}\neq\left\langle{\tiny\textsc{COPY}}\right\rangle$ we get $\Psi(\left\langle{\tiny\textsc{COPY}}\right\rangle)\cdot g_{t-1,i}=0$ and therefore $g_{t,i}^{*}=0$.

For now, assume that all the $n$-grams in ${\bm{x}}$ are unique, and that the length of the input satisfies $2L+2\leq K$ for $K=D^{n}$.

Proof: We separate to the following cases:

𝑖𝑛1\displaystyle=\bm{1}\{j>n\}\cdot[\Psi(x_{j-1}),\dots,\Psi(x_{j-n})]^{\top}[\Psi(x_{i}),\dots,\Psi(x_{i-n+1})] $\displaystyle=\bm{1}\{j>n\}\cdot\sum_{t=1}^{n}\Psi(x_{j-t})\Psi(x_{i-t+1})$ Now, if $j=i-L+1$ then $x_{j-t}=x_{i-L+1-t}=x_{i-t+1}$ and since $j>n$ we get

Proof: Using the previous claim, togehter with Lemma C.3, we get that:

$s_{i-L+1}\geq\frac{1}{1+i\exp(-\tau/d)}\geq\frac{1}{1+K\exp(-\tau/d)}\geq\frac{1}{1+\epsilon/2}=1-\frac{\epsilon/2}{1+\epsilon/2}\geq 1-\epsilon$

Claim: Fix some $\epsilon\in(0,1)$ and some $i>L$. Then, for $\tau\geq d\ln(\frac{2K}{\epsilon})$, it holds that:

𝑖𝐿1italic-ϵ\left\lVert H(\Psi(x_{1}),\dots,\Psi(x_{i}))-\Psi(x_{i-L+1})\right\rVert\leq\epsilon Proof: Let ${\bm{s}}$ as defined in the previous claim. Then:

Claim: If $\tau\geq d\ln(8Kd)$, then for all $i>L$ we have $\Phi(H(\Psi(x_{1}),\dots,\Psi(x_{i})))=x_{i-L+1}$.

Proof: Denote ${\bm{y}}_{i}=H(\Psi(x_{1}),\dots,\Psi(x_{i}))$. First, using the previous claim, we observe that

𝑖𝐿1\displaystyle{\bm{y}}_{i}\cdot\Psi(x_{i-L+1}) $\displaystyle=({\bm{y}}_{i}-\Psi(x_{i-L+1}))\cdot\Psi(x_{i-L+1})+\left\lVert\Psi(x_{i-L+1})\right\rVert$ $\displaystyle\geq 1-\left\lVert{\bm{y}}_{i}-\Psi(x_{i-L+1})\right\rVert\geq 1-\frac{1}{4d}$ Next, observe that for all $j\neq i-L+1$ we have

C.3 Proof of Lemma 2.4

Fix some $i<j\in[L]$. Let $I:=\{i,\dots,i+n\}$ and $J:=\{j,\dots,j+n\}$. We first bound the probability of drawing some ${\bm{x}}$ s.t. ${\bm{x}}_{I}={\bm{x}}_{J}$. Note that there are $D^{\left\lvert I\cup J\right\rvert}$ choices for ${\bm{x}}_{I\cup J}$. We count the number of choices for ${\bm{x}}_{I\cup J}$ s.t. ${\bm{x}}_{I}={\bm{x}}_{J}$. Notice that in this case, ${\bm{x}}_{I\cup J}$ is determined by ${\bm{x}}_{I\setminus J}$, therefore there are $D^{\left\lvert I\setminus J\right\rvert}$ possible choices. We conclude that

Appendix D Proofs - Lower Bound

In this section, we prove Theorem 2.7. We begin by showing that, for every input, the output of the model in each iteration is a deterministic function of the state of the model after observing the input:

Let $x_{n^{\prime}+1},\dots,x_{n^{\prime}+n}$ be the outputs of $H_{u,r}$. We need to show that there exist functions $G_{1},\dots,G_{n}$ s.t. $H_{u,r}(x_{1},\dots,x_{n^{\prime}})=G(S_{n^{\prime}}(x_{1},\dots,x_{n}))$. We give the following recurive definition:

$G_{1}(s)=r(s)=R_{n^{\prime}}(x_{1},\dots,x_{n^{\prime}})=x_{n^{\prime}+1}$.

Given the previous Lemma, we bound the error of the model by comparing the number of possible states to the number of possible inputs.