Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Authors: Albert Gu
Tri Dao
Venue: COLM (2024)
Keywords: sequence model, deep learning, state space model, S4, Mamba
TL;DR: This paper introduces the Mamba architecture.

Abstract: Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5x higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

Published as a conference paper at COLM 2024
Linear-Time Sequence Modeling with Selective State Spaces
Albert Gu"
Tri Dao"
Machine Learning Department Carnegie Mellon University
Department of Computer Science Princeton University
aguêcs. cmu. edu
tri@tridao.nd mee
Abstract
Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMss have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modal- ities such as language. We identify that a key weakness of such models is their in- ability to perform content- based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weak- ness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5x higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-lenghh sequences. As general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.
1 Introduction
Foundation models (FMs), or large models pretrained on massive data then adapted for downstream tasks, have emerged as an effective paradigm in modern machine learning. The backbone of these FMs are often sequence models, operating on arbitrary sequences of inputs from wide variety of domains such language, images, speech, audio, time series, and genomics (Sutskever et al., 2014; Dosovitskiy et al., 2020; Oord et al., 2016; Brown et al., 2020; Ismail Fawaz et al., 2019; Poli et al., 2023). While this concept is agnostic to a particular choice of model architecture, modern FMs are predominantly based on a single type of sequence model: the Transformer (Vaswani et al., 2017) and its core attention layer. The efficacy of self-attention is attributed its ability route information densely within a context window, allowing it to model complex data. However, this property brings fundamental drawbacks: an inability to model anything outside of finite window, and quadratic scaling with respect to the window length. An enormous body of research has appeared on more efficient variants of attention to overcome these drawbacks (Tay et al., 2022), but often at the expense of the very properties that makes it effective. As of yet, none of these variants have been shown to be empirically effective at scale across domains.
Recently, structured state space sequence models (SSMs) (Gu et al., 2021; 2022a) have emerged as a promising class of architectures for sequence modeling. These models can be interpreted as a combination of recurrent neural networks (RNNs) and convolutional neural networks (CNNs), with inspiration from classical state space models (Kalman, 1960). This class of models can be computed very efficiently as either a recurrence or convolution, with linear or near-linear scaling in sequence length. Additionally, they have principled mechanisms for modeling long-range dependencies (Gu et al., 2020a) in certain data modalities, and have dominated benchmarks such as the Long Range Arena (Tay et al., 2021). Many flavors of SSMs (Gu et al., 2022a; Gupta et al.,
* AAlphabetical by first name.

---

Published as a conference paper at COLM 2024
2022a; Gu et al., 2022b; Li et al., 2023; Ma et al,, 2023; Smith et al., 2023; Orvieto et al., 2023) have been successful in domains involving continuous signal data such as audio and vision (Goel et al., 2022; Saon et al., 2023; Nguyen et al., 2022). However, they have been less effective at modeling discrete and information-dense data such as text.
We propose a new class of selective state space models, that improves on prior work on several axes to achieve the modeling power of Transformers while scaling linearly in sequence length.
Selection Mechanism. First, we identify a key limitation of prior models: the ability to efficiently select data in input-dependent manner (i.e. focus on or ignore particular inputs). Building on intuition based on important synthetic tasks such as selective copy and induction heads, we design simple selection mechanism by parameterizing the SSM parameters based on the input. This allows the model to filter out irrelevant information and remember relevant information indefinitely.
Hardware-aware Algorithm. This simple change poses a technical challenge for the computation of the model; in fact, all prior SSMs models must be time- and input-invariant in order to be computationally efficient. We overcome this with a hardware-aware algorithm that computes the model recurrently with a scan instead of convolution, but does not materialize the expanded state in order to avoid 1O access between different levels of the GPU memory hierarchy. The resulting implementation is faster than previous methods both in theory (scaling linearly in sequence length, compared to pseudo-linear for all convolution-baser SSMs) and on modern hardware (up to 3 faster on A100 GPUs).
We simplify prior deep sequence model architectures by combining the design of prior SSM architectures (Dao et al., 2023) with the MLP block of Transformers into single block, leading to a simple and homogenous architecture design (Mamba) incorporating selective state spaces.
Architecture.
Selective SSMs, and by extension the Mamba architecture, are fully recurrent models with key prop- erties that make them suitable as the backbone of general foundation models operating on sequences. (i) High quality: selectivity brings strong performance on dense modalities such as language and ge- nomics. (ii) Fast training and inference: computation and memory scales linearly in sequence length during training, and unrolling the model autoregressively during inference requires only constant time per step since it does not require a cache of previous elements. (iii) Long context: the quality and efficiency together yield performance improvements on real data up to sequence length 1M.
We empirically validate Mamba's potential as general sequence FM backbone, in both pretraining quality and domain-specific task performance, on several types of modalities and settings:
Synthetics. On important synthetic tasks such as copying and induction heads that have been proposed as being key to large language models, Mamba not only solves them easily but can extrapolate solutions indefinitely long (> 1M tokens).
Audio and Genomics. Mamba out-performs prior state-of-the-art models such SaShiMi, Hyena, and Transformers on modeling audio waveforms and DNA sequences, both in pretraining quality and downstream metrics (e.g- reducing FID on a challenging speech generation dataset by more than half). In both settings, its performance improves with longer context up to million-lengthh sequences.
Language Modeling. Mamba is the first linear-time sequence model that truly achtieves Transformer- quality performance, both in pretraining perplexity and downstream evaluations. With scaling laws up to 1B parameters, we show that Mamba exceeds the performance of large range baselines, including very strong modern Transformer training recipes based on LLaMa (Touvron et al., 2023). Our Mamba language model has 5x generation throughput compared to Transformers of similar size, and Mamba-3B's quality matches that of Transformers twice its size (e.g. 4 points higher avg. on common sense reasoning compared Pythia-3B and even exceeding Pythia-7B)).
2 State Space Models
Structured state space sequence models (S4) are a recent class of sequence models for deep learning that are broadly related to RNNs, and CNNs, and classical state space models. They are inspired by particular continuous system (1) that maps 1-dimensional function or sequence x(t) E R - y(t) R through an implicit latent state It(1) E VR Concretely, S4 models are defined with four parameters (A,,,B,C), which define seguence-to-sequeee transformation in two stages.
2

---

Published as a conference paper at COLM 2024
((()= Ah()+++())+0)))) +Bx(I) (1a)
11= 1 + Bx
(2a)
K KCB,CAAB ...., CA CA...... B,...)
(3a)
Ch()=)(()
(Ib)
y1 -h=
(2b)
y=x*R
(3b)
Discretization. The first stage transforms the " "continuous parameters" (4, to "discrete parameters" () through fixed formulas A = fa (A,A) and B = fa(fA,((,) (A,,,, where the pair (fxff) is called a discretization rule. The most common is zero-order hold (ZOH) defined by A A=exp(AA) and B- (AA) (exp(AA)) -)B.
Discretization has deep connections to continuous-time systems which can endow them with additional properties such as resolution invariance (Nguyen et al,, 2022) and automatically ensuring that the model is properly normalized (Gu et al., 2023; Orvieto et al., 2023), It also has connections to gating mechanisms of RNNs (Tallec & Ollivier, 2018; Gu et al., 2020b) which we will revisit in Section 3.5. However, from mechanical point of view discretization can simply be viewed as the first step of the computation graph in the forward pass of an SSM
Computation. After the parameters have been transformed from (A,A,B,C) (A,B,C), the model can be computed in two ways, either as linear recurrence (2) or a global convolution (3).
Commonly, the model uses the convolutional mode (3) for efficient parallelizable training (where the whole input sequence is seen ahead of time), and switched into recurrent mode (2) for efficient autoregressive inference (where the inputs are seen one timestep at time).
Linear Time Invariance (LTI). An important property of equations (1) to (3) is that the model's dynamics are constant through time. In other words (A,,, and consequently (A,B) well, are fixed for all time-steps. This property is called linear time invariance (LTI), which is deeply connected to recurrence and convolutions. Informally, we think of LTI SSMs being equivalent to any linear recurrence (2a) or convolution (3b), and use LTI as an umbrella term for these classes of models.
Thus far, all structured SSMs have been LTI (e.g. computed as convolutions) because of fundamental efficiency constraints, discussed in Section 3.3. However, a core insight of this work is that LTI models have fundamental limitations in modeling certain types of data, and our technical contributions involve removing the LTI constraint while overcoming the efficiency bottlenecks.
Structure and Dimensions. Finally, we note that structured SSMs are so named because computing them efficiently also requires imposing structure on the A matrix. The most popular form of structure is diagonal (Gupta et al., 2022a; Gu et al., 2022b; Smith et al., 2023), which we also use.
In this case, the A AER' NxN B Nx1 ,ee RR matrices can all be represented by N numbers. To operate over an input sequence x of batch size B and length L with D channels, the SSM is applied independently to each channel. Note that in this case, the total hidden state has dimension DN per input, and computing it over the sequence length requires O(BLDN) time and memory; this is the root of the fundamental efficiency bottleneck addressed in Section 3.3.
SSM Architectures. SSMs are standalone sequence transformations that can be incorporated into end-to-end neural network architectures. We discuss some of the most well-known SSM architectures, many of which will also serve as our primary baselines.
Linear attention (Katharopoulos et al., 2020) is an approximation of self-attention involving recurrence which can be viewed as a degenerate linear SSM.
H3 (Dao et al., 2023) generalized this recurrence to use S4; it can be viewed as an architecture with an SSM sandwiched by two gated connections (Figure 2). H3 also inserts standard local convolution, which they frame as a shift-SSM, before the main SSM layer.
Hyena (Poli et al., 2023) uses the same architecture as H3 but replaces the S4 layer with an MLP- Lull-paramtterized global convolution (Romero et al., 2021).
Other closely related SSMs and architectures are discussed further in an extended related work (Appendix B). We highlight in particular $5 (Smith et al., 2023), QRNN (Bradbury et al., 2016), and SRU (Lei et al., 2017), which we view the most closely related methods to our core selective SSM.
3 Selective State Space Models
We motivate our selection mechanism using intuition from synthetic tasks (Section 3.1), then explain how to incorporate this mechanism into state space models (Section 3.2). The resulting time-varying SSMs cannot use convolutions, presenting technical challenge of how to compute
3

---

Published as a conference paper at COLM 2024
Copying
Selective Copying
Output NARNNRN ...
Outoul N N
Input
Input
Solution
Induction Heads
Pedletsy
Figure (Leff) The standard version of the Copying task involves constant spacing between input and output el- ements and is easily solved by time-invariant models such as linear recurrences and global convolutions. (Right Top) The Selective Copying task has random spacing in between inputs and requires time-varying models that can selectively remember or ignore inputs depending on their content. (Right Botfom) The Induction Heads task is example of associative recall that requires retrieving answer based on context, key ability for L.Mss.
them efficiently. We overcome this with hardware-aware algorithm that exploits the memory hierarchy on modern hardware (Section 3.3). We then describe a simple SSM architecture without attention or even MLP blocks (Section 3.4). Finally, we discuss some additional properties of selection mechanisms (Section 3.5).
3.1 Motivation: Selection Means of Compression
We argue that fundamental problem of sequence modeling is compressing confext into a smaller stafe. In fact, we can view the tradeoffs of popular sequence models from this point of view. For example, attention is both effective and inefficient because it explicitly does not compress context at all. This can be seen from the fact that autoregressive inference requires explicitly storing the entire context (i.. the KV cache), which directly causes the slow linear-time inference and quadratic-tinn time training of Transformers. On the other hand, recurrent models are efficient because they have finite state, implying constant-time inference and linear-time training. However, their effectiveness is limited by how well this state has compressed the context.
To understand this principle, we focus two running examples of synthetic tasks (Figure 1).
The Selective Copying task modifies the popular Copying task (Arjovsky et al., 2016) by varying the position of the tokens to memorize. It requires content-mware reasoning to be able to memorize the relevant tokens (colored) and filter out the irrelevant ones (white).
The Induction Heads task is a well-known mechanism hypothesized to explain the majority of in-context learning abilities LLMs (Olsson et al., 2022). It requires context-mone mwre reasoning to know when to produce the correct output in the appropriate context (black).
These tasks reveal the failure mode of LTI models. From the recurrent view, their constant dynamics (e.g. the (A,B) transitions in (2)) cannot let them select the correct information from their context, or affect the hidden state passed along the sequence in an input-dependent way. From the convolutional view, it is known that global convolutions can solve the vanilla Copying task (Romero et al., 2021) because it only requires time-awareness, but that they have difficulty with the Selective Copying task because of lack of content-awareness (Figure 1). More concretely, the spacing between inputs-to-outputs is varying and cannot be modeled by static convolution kemels.
In summary, the efficiency vs. effectiveness tradeoff of sequence models is characterized by how well they compress their state: efficient models must have a small state, while effective models must have a state that contains all necessary information from the context. In turn, we propose that a fun- damental principle for building sequence models is selectivity: or the context-aware ability to focus on or filter out inputs into a sequential state. In particular, selection mechanism controls how infor- mation propagates or interacts along the sequence dimension (see Section 3.5 for more discussion).
3.2 Improving SSMs with Selection
One method of incorporating a selection mechanism into models is by letting their parameters that affect interactions along the sequence (e.g. the recurrent dynamics of an RNN or the convolution kernel of a CNN) be input-dependent.
Algorithms 1 and illustrates the main selection mechanism that we use. The main difference is simply making several parameters A,B,C functions of the input, along with the associated changes to tensor shapes throughout. In particular, we highlight that these parameters now have a length dimension I., meaning that the model has changed from time-invariant to time-varying. (Note
4

---

Published as a conference paper at COLM 2024
Algorithm SSM (S4)
Algorithm 2 SSM Selection (S6)
Input: x: (B,,LD))
Input: (:,,L,
Output: y: (B,LD)
Oulput: y: (B,L,)
1: (D(,, Parameter
1: A: (D,N) < Parameter
> Represents structured N x N matrix
Represents structured N X matrix
2: (D(,,N) <- Parameter
2: B: (B,L,, <- SB(x)
3: (D,N) Parameter
3: (B,L,N)) f- sc(x)
4: (D( To (Parameter)
(B,L,, < TA (Parameter+-sg
5: A,B: (D,0) <- discretizeeA,AA,
5: A,B: (B,,,D, DiscretizeeA,A,,B
6: y+ (V-SSM(A,B,C)(a)
6: + SSM((A,,.)(x)
b Time-invariant recurrence or convolution
> Time-varying: recurrence (scan) only
7: returny y
7: returny y
that shape annotations were described in Section 2.) This loses the equivalence to convolutions (3) with implications for its efficiency, discussed next.
We specifically choose $B (x) Linear sc(x) = Lineary 83 (x) Broadcastp (Lineary (x)), and Ta = softplus, where Linear, is a parameterized projection to dimension d. The choice of Ss and Ta due to connection to RNN gating mechanisms explained in Section 3.5.
3.3 Efficient Implementation of Selective SSMs
Hardware-friendll primitives such as convolutions (Krizhevsky al., 2012) and attention (Bah- danau et al., 2015; Vaswani et al., 2017) enjoy widespread application. Here we aim to make selective SSMs efficient modern hardware (GPUs) as well. The selection mechanism is quite natural, and earlier works attempted to incorporate special cases of selection, such as letting A vary over time in recurrent SSMs (Gu et al., 2020a). However, this was computationally difficult, which was why and derivatives used LTI (non-selective) models, most commonly in the form of global convolutions.
The selection mechanism is designed to overcome the limitations of LTI models; at the same time, we therefore need to revisit the computation problem of SSMs. We address this with three classical techniques: kernel fusion, parallel scan, and recomputation. We make two main observations:
The naive recurrent computation uses O(BDDN) FLOPs while the convolutional computation uses O O(BLDlog(L) (L) FLOPs, and the former has lower constant factor. Thus for long sequences and not-too-large state dimension N, the recurrent mode can actually use fewer FLOPs.
The two challenges are the sequential nature of recurrence, and the large memory usage. To address the latter, just like the convolutional mode, we can attempt to not actually materialize the full state
The main idea is to leverage properties of modern accelerators (GPUs) to materialize the state h only in more efficient levels of the memory hierarchy. In particular, most operations (except matrix multiplication) are bounded by memory bandwidth (Williams et al., 2009; Ivanov et al., 2021; Dao et al., 2022). This includes our scan operation, and we use kernel fusion to reduce the amount of memory IOs, leading significant speedup compared
standard implementation.
Concretely, instead of preparing the scan input (A, B) of size (B,LD,N) in GPU HBM (high- bandwidth memory), we load the SSM parameters (A,A,B,C) directly from slow HBM fast SRAM, perform the discretization and recurrence in SRAM, and then write the final outputs of size (B,LD) back HBM.
To avoid the sequential recurrence, we observe that despite not being linear it can still be parallelized with a work-efficient parallel scan algorithm (Blelloch, 1990; Martin & Cundy, 2018; Smith et al., 2023).
Details of the fused kernel and recomputation are in Appendix D.
3.4 Simplified SSM Architecture
As with structured SSMs, selective SSMs are standalone sequence transformations that can be flexibly incorporated into neural networks. The H3 architecture is the basis for the most well-known SSM architectures (Section 2), which are generally comprised of a block inspired by linear attention interleaved with an MLP (multi-layer perceptron) block. We simplify this architecture by combining these two components into one, which is stacked homogenously (Figure 2).
5

---

Published as a conference paper at COLM 2024
Linear projection
SSM
SSM
Sequancee
martoormatio
Conv
Nent nearty mulitiiicttin
H3
Gated MLP
Mamba
Figure 2: (Architecture.) Our simplified block design combines the H3 block, which is the basis of most SSM architectures, with the ubiquitous MLP block of modemn neural networks. Instead of interleaving these two blocks, we simply repeat the Mamba block homogenously, Compared to the H3 block, Mamba replaces the first multiplicative gale with an activation function. Compared the MLP block, Mamba adds an SSM the main branch. For i we use the SiL.U/ Swish activation (Hendrycks Gimpel, 2016; Ramachandran et al., 2017).
This architecture involves expanding the model dimension D by controllable expansion factor F. For each block, most of the parameters (3ED²) are in the linear projections while the inner SSM contributes less. We always fix to E = 2 in our experiments and use two stacks of the block to match the 12D2 parameters of a Transformer's 's interleaved MHA (multi-head attention) and blocks.
3.5 Properties of Selection Mechanisms
The selection mechanism is a broader concept that can be applied in different ways, such as to more traditional RNNs or CNNs, to different parameters (e.g. A in Algorithm 2), or using different transformations s(x).
We highlight the most important connection: the classical gating mechanism of RNNs is an instance of our selection mechanism for SSMs. We note that the connection between RNN gating and the discretization of continuous time systems is well established (Funahashi & Nakamura, 1993; Tallec & Ollivier, 2018). In fact, Theorem 1 is an improvement of Gu et al. (2021, Lemma 3.1) generalizing to the ZOH discretization and input-dependent gates (proof in Appendix C). More broadly, A in SSMs can be seen to play generalized role of the RNN gating mechanism. In line with prior work, we adopt the view that discretization of SSMs the principled foundation of heuristic gating mechanisms.
Theorem 1. When N = 1,A = 1,B = 1, Linear(x), and Ta softplus, then the selective SSM recurrence (Algorithm 2) takes the form = (x))ear((x)) (the gate) and h4 = (1 g: -
1 g:
As mentioned in Section 3.2, our specific choices SA,TA is from this connection. In particular, note that if a given input x; should be completely ignored (as necessary in the synthetic tasks), all channels should ignore it, and so we project the input down to 1 dimension before repeating/ httgg//braaaaiinggg with A.
We elaborate on two particular mechanistic effects of selection.
Variable Spacing. Selectivity allows filtering out irrelevant noise tokens that may occur between inputs of interest. This is exemplified by the Selective Copying task, but occurs ubiquitously in common data modalities, particularly for discrete data - for example the presence of language fillers such "um". This property arises because the model can mechanistically filter out any particular imput xt, for example in the gated RNN case (Theorem 1) when gt +
Filtering Context. has been empirically observed that many sequence models do not improve with longer context (Shi et al., 2023a), despite the principle that more context should lead to strictly better performance. An explanation is that many sequence models cannot effectively ignore irrelevant context when necessary; an intuitive example are global convolutions (and general LTI models). On the other hand, selective models can simply reset their state at any time to remove extraneous history, and thus their performance in principle improves monotonicly with context length (e.g. Appendix F.2.2).
We remark that while the A parameter could also be selective, it ultimately affects the model only through its interaction with A via A = ex)(AAA) (the discretization. Thus selectivity in A is enough to ensure selectivity in (A, and is the main source of improvement. We hypothesize that making A selective in addition to (or instead of) A would have similar performance, and leave it out for simplicity.

---

Published as a conference paper at COLM 2024
MODEL ARCH.
LAYER ACC.
Induction Heads Extrapolation
1.0-
S4
No gate S4
18.3
No gate
S6
97.0
0.8
MHA Absolute
MHA--R
H3
H3
S4
57.0
Accuracy
0.6
MHA
H3
Hyena
H3
Hyena
30.1
Hyena
H3
S6
99.7
0.4
Mamba
Random
Mamba S4
56.4
0.2
Langth
Mamba Hyena
28.4
0.0
Mamba
Mamba
S6
99.8
16a
10
10
107
10r
Test Sequence Length
Figure 3: (Induction Heads.) Models are trained on sequence length
Table 1: (Selective Copying,)
Accuracy for combinations of archi- tectures and inner sequence layers.
2 256, and tested on increasing sequence lengths of 2 64 up to 2 = 1048576. Full numbers in Table 10.
Scaling (Sequence Length 2048)
Scaling Laws (Sequence Length 0192)
2'-13
Perpleaity (big stal)
5MRY
Perplesity (bg (oea)
540CY
Trandoreer Ratten
Travsfareer
Redfet
Trando
10
FLOPs (log soale)
FLOPs 0og scale)
Figure 4: (Scaling Laws.) Models size 125M to Rs 1.3B parameters, trained on the Pile. Mamba scales better than all other attention-free models and is the first to match the performance of a very strong "Transformer=+" recipe that has now become standard, particularly as the sequence length grows.
Remark 3.1. For brevity in our experimental results, we sometimes abbreviate selective SSMs as S6 models, because they are S4 models with a selection mechanism and computed with a scan.
Empirical Evaluation
Mamba achieves thet---rtthh-tttte results on the synthetic tasks (Section 4.1) and three different domains (language, DNA, audio) (Sections 4.2 to 4.4) on both pretraining and downstream tasks, while being very computationally efficient (Appendix E.4).
4.1 Synthetic Tasks
Table and Figure 3 show results for the synthetic tasks. On Selective Copying, the selective SSM layer is enough to solve the task independently of the architecture used, while previous LTT SSMs cannot even when combined with more powerful architectures. On Induction Heads, Mamba learns the task perfectly and can even extrapolate to million-length sequences, or 4000 longer than it saw during training, while no other method goes beyond 2x. Full discussion for synthetic tasks are in Appendix E.1.
4.2 Language Modeling
We evaluate the Mamba architecture on standard autoregressive language modeling against other architectures, on both pretraining metrics (perplexity) and zero-shot evaluations. We set the model sizes (depth and width) to mirror GPT3 specifications. We use the Pile dataset (Gao et al., 2020), and follow the training recipe described in Brown et al. (2020). All training details are in Appendix F.2.
Scaling Laws. For baselines, architecture), as well as the strongest Transformer recipe we know of (here referred to as Transformer++4), based on the PaLM and LLaMa architectures (e.g, rotary embedding, SwiGLU MLP, etc.). We also compare against other recent subquadratic architectures (Figure 4). All model details are in Appendix F.2.
compare against the standard Transformer architecture (GPT3
Downstream Evaluations. Table 2 shows the performance of Mamba on range of popular down- stream zero-shot evaluation tasks. We compare against the most well-known open source models at these sizes, most importantly Pythia (Biderman et al., 2023) and RWKV (Peng et al., 2023) which were trained with the same tokenizer, dataset, and training length (300B tokens) as our models.
7

---

Published as a conference paper at COLM 2024
Table 2: (Zero-shot Evaluations.) Best results for each size in bold. We compare against open source LMs with various tokenizers, trained for up to 300B tokens. Pile refers the validation split, comparing only against models trained on the same dataset and tokenizer (GPT-NeeX-20D5) For each model size, Mamba best-in-class on every single evaluation result, and generally matches baselines at twice the model size.
MODEL
TOKEN.
PILE
LAMBADA LAMBADA
HELLASWAG
PIQA ACC
ARC-E
ARC-C WINOGRANDE
AVERAGE
PPL
PFI.
ACC
ACC
ACCT T
ACC 7
ACC1
ACC
Hybrid HB-30IM
GPT2
12.58
48.0
41.5
68.1
51.4
24.7
54.1
48.0
Pythia-IOM
NeoX
9.95
10.84
514
40.6
66.9
52.1
24.6
5388
462
Mamba-370M
NeoX
8.28
8.14
55.6
46.5
69.5
28.0
55.3
50.0
Pythia-IB
NeoX
7,82
7.92
56.1
47.2
70.7
27.1
53.5
51.9
Mamba-79UM
NeoX
7.33
6.02
62.7
55.1
72.1
61.2
29.5
56.1
57.1
GIT--Noo 1.38
GPT2
7.50
71.1
54.9
Hybrid H3-1.38
GPT2
11.25
49.6
52.6
71.3
39.2
28. 1
56.9
53.0
OIT--1.3B
72.4
Pythia-1.4B
NeuX
751
6.08
61.7
521
71.0
60.5
28.5
57.2
55.2
RWKKV-1.5B
NeoX
7.70
7.04
56.4
52.5
72.4
54.6
54.3
Mamba-1.4B
NeoX
6.80
5.04
64.9
59:1
74.2
65.5
32.8
61.5
59.7
GPT-Neo
5.63
72.1
Hybrid H3-278
GPT2
7.92
55.7
59.7
73.3
65.6
32.3
61.4
58.0
OPT--27B
OPT
5.12
74.8
58.7
Pythia-2.88
NeoX
6.73
5.04
64.7
59.3
74.0
64.1
32,9
59.7
59.1
RWKKV-3B
NeuX
7.00
5.24
73.7
59.6
59.6
Mamba-2.88
NeoX
6.22
4.23
69.2
66.1
75.2
69.7
36.3
63.5
63.3
GPT-1-6B
GPT2
4.10
68.3
66.3
754
67.0
36.6
64.1
63.0
OPT-6.7B
OPT
4.25
67.7
67.2
76.3
65.6
34.9
65.5
62.9
Pythia-6.9B
NeuX
6.51
4.45
75.2
61.3
RWKKV77.BB
NeoX
6.31
4.38
67.2
65.5
76.1
67.8
37.5
61.0
62.5
Scaling
(H638)
Scaling Laws Sequence Length (H638)
3.00
31
Mania Tranforment
Mavita
Membe
Perpleatt
Perpleaty
2.50
2:00
28-
277
Parameters Olgg
Seguence Langth
Figure 3: (DNA Scaling Laws.) Pretraining on the HG38 (human genome) dataset. (Left) Fixing short context length 20 1024 and increasing size from 200K 40M parameters, Mamba scales better than baselines. (Right) Fixing model size and increasing sequence lengths while controlling for computation. Unlike baselines, the selection mechanism of Mamba facilitates better performance with increasing context length.
4.3 DNA Modeling
Motivated by the success of large language models, there has been recent exploration into using the foundation model paradigm for genomics. DNA has been likened to language in that it consists of sequences of discrete tokens with a finite vocabulary. long-range dependencies to model (Avsec et al., 2021). We investigate Mamba as FM backbone for pretraining and fine-tuning in the same setting recent works long-sequence models for DNA (Nguyen et al., 2023). In particular, we focus on two explorations of scaling laws across model size and sequence length (Figure 5), and difficult downstream synthetic classification task requiring long context (Figure 6). Full discussion for these tasks is in Appendix E.2.
also known for requiring
4.4 Audio Modeling and Generation
For the audio waveform modality, we compare primarily to the SaShiMi architecture and training protocols (Goel al., 2022). The architecture is U-Net with alternating S4 and MLP blocks, which we consider replacing with Mamba. Experiment details are in Appendix F44
Long-Context Autoregressive Pretraining. We evaluate pretraining quality on YouTubeMix (Deep- Sound, 2017), standard piano music dataset. Figure/ evaaluates the effect of increasing training sequence lengths from 2 -812 2 109, while keeping computation fixed.
Autoregressive Speech Generation. Tables 3 and show results on the SC09 dataset for generating speech clips of the digits 019 (Warden, 2018; Donahue et al., 2019). Mamba sets significant state-of-the. art results, and is consistently better than attention S4 blocks in parameter-contron setting.
8

---

Published as a conference paper at COLM 2024
Accuracy Especiso
Scaling Laws -Squeene i Length (YouTibeMla)
1.473
Mantta 1.M
1.400
Mariaa Fandom
1.42
Accunooy
1.800
1:375
1.35
1383
64-
1300
Length
Langth
Figure 6: (Great Apes DNA Classification.) Accuracy after fine-tuning on sequences of length 2 = 1024 up to 2 1048576 using pretrained models the same context length. Numerical results in Table 12.
Figure 7: (Audio Pretraining.) Mamba improves performance prior state-of-the-art (Sashimi) in autoregressive audio modeling, while improving up to minute-long context or million-length sequences.
Table 3: (SC09) Automated metrics for unconditional Table 4: (SC09 Model Ablations) Models with 6M generation on challenging dataset of fixed-length parameters. In SaShiMi's U-Net backbone, there speech clips. (Top to Boltom) Autoregressive baselines, non-autoregressive baselines, Mamba, and dataset 1000, sandwiched metrics.
center blocks operating sequence length
each side by 8 outer blocks on sequence length 4000, sandwiched by 8 outer blocks on sequence length 16000 (40 blocks total). The architecture of the center blocks are ablated independently of the rest. Note that Transformers (MHA MMLP) were not tested in the more important outer blocks because of efficiency constraints.
MODEL
PARAMS
NLL.L
FID J
IS
MIS
AM 1
SampleRNN
35.0M
2,022
8.96
1.71
3.02
1.76
WaveNet
42M
1925
5.08
2.27
1447
SaShiMi
58M
1.873
1.99
5.13
42.57
0.74
WaveGAN
19.1M
2.03
4.90
36.10
0.80
DiftWave
OUTER
NLL
FID
MIS 1
+ SaShiMii
23.0M
1.42
5.94
69.17
0.59
S4-MEP
0.70
Mamba
611M
1,852
0.94
6.26
89.54
0.52
S4-MLP
SH+MLP
1.867
1.43
5.42
53.54
0.65
Mamba
24.3M
1.860
0.67
7.33
144.9
0.36
S4-MEP
56.51
Mamba
MHA-MLP
1.850
1.37
5.63
58
0.62
Train
0.00
8.56
292.5
0.16
Mamba
73.34
0.55
Test
0.02
8.33
257.6
0.19
Mamba
Mamba
1882
0.94
6.26
88.54
0.52
Scan us Convolation va Attention time (A100 80C8S PCIe)
Inference thrngghast 80GB jaramg leugh 2248)
Rail #
Thaotween
100
(unn))
Throughos Tobkers
Treufhreet 5.3
D0M
10
612
32c
64k
199c
Sequence longth
Figure 8: (Efficiency Benchmarks.) (Left) Training: efficient scan is 40x faster than standard implemen- tation. (Right) Inference: recurrent model, Mamba can achieve 5x higher throughput than Transformers.
Conclusion
We introduce selection mechanism to structured state space models, allowing them to perform context-dependent reasoning while scaling linearly in sequence length. When incorporated into a simple attention-free architecture, Mamba achieves state-of-the-al results on diverse set of domains, where it matches exceeds the performance of strong Transformer models. We are excited about the broad applications of selective state space models to build foundation models for different domains, especially in emerging modalities requiring long context such genomics, audio, and video. Our results suggest that Mamba is a strong candidate be a general sequence model backbone.
Acknowledgments
We thank Karan Goel, Arjun Desai, and Kush Bhatia for helpful feedback the draft.

---

Published as a conference paper at COLM 2024
References
Martin Arjovsky, Amar Shah, and Yoshua Bengio, Unitary evolution recurrent neural networks. In The International Conference on Machine Learning (ICML), pp. 1120-1128, 2016.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Strbska-Barwwinse Kyle Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature Methods, 18(10):119-1103, 2021.
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin lonescu. Using fast weights to attend to the recent past. Advances in Neural Information Processing Systems (NeurIPS), 29, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly leamning to align and translate. In The International Conference on Learning Representations (ICLR), 2015.
David Balduzzi and Muhammad Ghifary. Strongly-typed recurrent neural networks. In International Conference Machine Learning, pp. 1292-1300. PMLR, 2016.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In The International Conference on Machine Learning (ICML), pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on Artificial Intelligence, volume 34, 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et al. Gpt-neox-2066 An open-source autoregressive language model. arXiv preprint arXiv:2204.06745 2022.
Guy E Blelloch. Prefix sums and their applications. 1990.
James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. arXiv preprint ar Xie:16611.0157,, 2016.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Procssing Systems (NeurIPS), 33:1877-1901, 2020.
Aydar Bulatov, Yuri Kuratov, and Mikhail S Burtsev. Scaling Transformer to 1M tokens and beyond with RMT. arXiv preprint ar arxii:2304.10666 2023.
Rewon Scott Transformers. arXiv preprint ar arXiv:1904..00 10509, 2019.
Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In The International Conference on Learning Representations (ICLR), 2021.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. PaLM: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023. URL http:t/m////m httpiitrr..aaaaaaeeeeeeeee wttgg//prrs///eeee /224/222222222// 1144. html.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio, Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint
2014.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the A12 reasoning challenge. arXi preprint ar Xiv: 103.0055,, 2018
10

---

Published as a conference paper at COLM 2024
Tri Dao. FlashAttention--:: Faster attention with better parallelism and work partitioning. In The International Conference on Learning Representations (ICLR), 2024.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advmnces in Neural Information Processing Systems (NeurIPS), 2022.
Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The International Conference on Learning Representations (ICLR), 2023.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In The International Conference on Machine Learning (ICML), pp. 933-941. PMLR, 2017.
DeepSound. Samplemn. http:://ttttttooooo/ooooo//////r
2017.
Jiayu Ding, Shumiing Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, and Furu Wei. LongNet: Scaling Transformers to 1,000,000,000 tokens. arXiv preprint arXio:2307.024866 2023,
Chris Donahue, Julian McAuley, and Miller Puckette. Adversarial audio synthesis. In The International Conference on Learning Representations (ICLR), 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. In The International Conference on Learning Representations (ICLR), 2020.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac HatfieldDDodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for Transformer circuits. Transformer Circuits Thread, 2021.
/ttplisspppb///////aaaooo/////// /ttr///mwor/iiieee...eea /index.ttll...
2021
Mahan Fathi, Jonathan Pilault, Pierre-Luc Bacon, Christopher Pal, Orhan Firat, and Ross Goroshin. Block-State Transformer. ar arviv preprint ar Xiv::330.009539, 2023.
Yassir Fathullah, Chunyang Wu, Yuan Shangguan, Junteng Jia, Wenhan Xiong, Jay Mahadeokar, Chunxi Liu, Yangyang Shi, Ozlem Kalinli, Mike Seltzer, and Mark J. F. Gales. Multi-Head State Space Model for Speech Recognition. In Proc. INTERSPEECH 2023, pp. 241-245, 2023. doi: 10.21137/ (011437/Intern http://eeesseechh22220000000
Daniel Y Fu, Elliot L Epstein, Eric Nguyen, Armin W Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling. The International Conference Machine Learning (ICML), 2023.
Ken-ichi Funahashi and Yuichi Nakamura. Approximation of dynamical systems by continuous time recurrent neural networks. Neural Networks, 66))8001-06, 1993.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800GB dataset of diverse text for language modeling. arviv preprint axXiv:2101.0002 0020100000000,,.. 2020.
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, September 2021. URL lattps://don http:///ooooo////////// 10. 5281/zenodo 5371628.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. In The International Conference on Machine Learning (ICML), 2022.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HIPPO: Recurrent memory with optimal polynomial projections. In Advances in Newral Information Processing Systems (NeurIPS), 2020a.
11

---

Published as a conference paper at COLM 2024
Albert Gu, Caglar Gulcehre, Tom Le Paine, Matt Hoffman, and Razvan Pascanu. Improving the gating mechanism recurrent neural networks. The International Conference on Machine Learning (ICML), 2020b.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-ting models with the linear state space layer. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems (NeurlPS), 2022b.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your HIPPO: State space models with generalized basis projections. In The International Conference on Learning Representations (ICLR), 2023.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Adoances in Neural Information Processing Systems, 35:22982-22999 2022a.
Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear RNNs. arXiv preprint axXiv:2212.00766 2022b.
David Ha, Andrew Dai, and Quoc V. Le. HyperNetworks. In The International Conference on Learning Representations (ICLR), 2017.
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. In The International Conference on Learning Representations (ICLR), 2023.
Mikael Henaff, Arthur Szlam, and Yann LeCun. Recurrent orthogonal networks and long-memory tasks. In The International Conference on Machine Learning (ICML), 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv preprint
arXio:1606.n8845, 2016.
Sepp Hochreiter.
Untersuchungen dynamischen neuronalen netzen. Diploma, Technische Universität München, 91(1):31, 1991.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8):
1735-1780, 1997.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: The difficulty of learning long-term dependencies, 2001.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et all An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems (NeuriPS), 35:30016-30030, 2022.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In The International Conference on Machine Learning (ICML), pp. 9099-9117. PMLR, 2022.
Hassan Ismail Fawaz, Germain Forestier, Jonathan Weber, Lhassane Idoumghar, and Pierre- Alaiie-Alain Muller. Deep learning for time series classification: A review. Data Mining and Knowledge Discovery, 33(4):917-963 2019.
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing Transformers. Proceedings of Machine Learning and Systems, 3:711-732, 2021.
Li Jing, Caglar Gulcehre, John Peurifoy, Yichen Shen, Max Tegmark, Marin Soljacic, and Yoshua Bengio. Gated orthogonal recurrent units: On learning to forget. Neural Computation, (4): 765-783, 2019.
12

---

Published as a conference paper at COLM 2024
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNNs: Fast autoregressive Transformers with linear attention. In International Conference Machine Learning, pp. 5156-5165. PMLR, 2020.
Shiva Kaul. Linear dynamical systems as a core computational primitive. Advanccss in Neural Information Processing Systems, 33:16808-1682) 2020.
Chrysoula Kosma, Giannis Nikolentzos, and Michalis Vazirgiannis. Time-parameterizedd convolu- tional neural networks for irregularly sampled time series. arXiv preprint arXiv:2308. 2023.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. ImageNet classification with deep con- volutional neural networks, Advances in Neural Information Processing Systems (NeurIPS), 25, 2012.
Tao Lei. When attention meets fast recurrence: Training language models with reduced compute. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7633-7648, 2021.
Tao Lei, Yu Zhang, Sida I Wang, Hui Dai, and Yoav Artzi Simple recurrent units for highly parallelizable recurrence. arXiv preprint arXi::1709002555, 2017.
Mario Lezcano-Casade and David Martinez- Cheap orthogonal constraints in neural networks: A simple parametrization the orthogonal and unitary group. In The International Conference on Machine Learning (ICML), 2019.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? In The International Conference on Learning Representations (ICLR), 2023.
Vasileios Lioutas and Yuhong Guo. Time-aware large kernel convolutions. In The International Conference on Machtne Learning (ICML), pp. 6172-6183. PMLR, 2020.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Shahar Lutati, Itamar Zimerman, and Lior Wolf. Focus your attention (with adaptive iir filters). arXio preprint ar arXiv:2305.14952 2023.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The International Conference on Learning Representations (ICLR), 2023.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In The International Conference on Learning Representations (ICLR), 2018.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language mod- eling via gated state spaces. In The International Conference on Learning Representations (ICLR), 2023.
Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey. Efficient orthogonal parametrisation of recurrent neural networks using Householder reflections. In International Conference on Machine Learning, pp. 2401-2409. PMLR, 2017.
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4ND: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Eric Nguyen, Michael Poli, Marjan Faizi, Armin Thomas, Callum Birch-Sykes, Michael Womow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, al. HyenaDNA: Long-range genomic sequence modeling at single nucleotide resolution. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
13

---

Published as a conference paper at COLM 2024
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits 2022/ine centfformmmrrcccuussss.p... /202 nn-contex--eamming-aae wenpeee--laaannnnnn--nnnd-ed -ndduuttionn----
Thread, 2022. https:/
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. WaveNet: A generative model 2016.
for raw audio. arXiv preprint ar
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In The International Conference on Machine Learning (ICM)), 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, pp. 1525-1534, 2016.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning, pp. 1310-1318, 2013.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. RWKV: Reinventing RNNs for the Transformer era. arxiv preprint ar arXtv:2305.1304 2023.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In The International Conference on Learning Representations (ICLR), 2021.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In The International Conference on Machine Learning (ICML), 2023.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. ar preprint ar xxXi:22100.00440 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. CosFormer: Rethinking softmax in attention. In The International Conference on Learning Representations (ICLR), 2022b.
Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The International Conference on Learning Representations (ICLR), 2023.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in Neural Information Processing Systems (NeurIPS), 20, 2007.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Swish: A self-gated activation function. arv preprint Xfo::77710005941, 7(1):5, 2017.
David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomczak, and Mark Hoogendoom. CK- Conv: Continuous kernel convolution for sequential data. Xiv preprint ar arXio:2102..0261 2021.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial Winograd schema challenge at scale. Communicationss of the ACM, 64(9):99-106, 2021.
George Saon, Ankit Gupta, and Xiaodong Cui. Diagonal state space augmented Transformers for speech recognition. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Imaool Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers are secretly fast weight programmers. In The International Conference on Machine Learning (ICML), pp. 9355-9366. PMLR, 2021.
14

---

Published as a conference paper at COLM 2024
Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-33, 1992.
Noam Shazeer. GLU variants improve Transformer. arXiv preprint arXi::2002.0 2020.
Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed H Chi, Nathanael Schärli, and Denny Zhou. Large language models can be easily distracted by irrelevant context. In The International Conference on Machine Learning (ICML), pp. 31210-31227. PMLR, 2023a.
Jiaxin Shi, Ke Alexander Wang, and Emily Fox. Sequence modeling with multiresolution convolutional memory. In The International Conference on Machine Learning (ICML), pp. 31312-31327. PMLR, 2023b.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. In The International Conference Learning Representations (ICLR), 2023.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced Transformer with rotary position embedding. ar preprint ar arXi::2104.09664 2021.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to Transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.
Ilya Sutskever, Oriol Vinyals, and Quoc Le. Sequence to sequence learning with neural networks. Advances in Neural Information Processing Systems (NeurIPS), 27, 2014.
Corentin Tallec and Yann Ollivier. Can recurrent neural networks warp time? In The International Conference on Learning Representations (ICLR), 2018.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long Range Arena: A benchmark for efficient Transformers. In International Conference on Learning Representations (ICLR), 2021.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient Transformers: A survey.
ACM Computing Surneys, 55(6):1-28, 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprini ar21122302.33971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Eugene Vorontsow, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. In International Conference on Machine Learning, pp. pp.3570-300 3570-3578. PMLR, 2017.
Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6387-6397, 2023.
Pete Warden. Speech commands: A dataset for limited-vocabularyy speech recognition. ArXiv, abs/1804.03209, 2018.
Samuel Williams, Andrew Waterman, and David Patterson.
Roofline: An insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65-6, 2009.
Brandon Yang, Gabriel Bender, Quoc V Le, and Jiquan Ngiam. CondConv: Conditionally
parameterized convolutions for efficient inference. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Compulational Linguistics, 2019.
15

---

Published as a conference paper at COLM 2024
Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free Transformer. arXio preprint arXiv:2105.14003 2021.
Lin Zheng, Chong Wang, and Lingpeng Kong. Linear complexity randomized self-attention mechanism. In International Conference on Machine Learning, pp. 27011-27041. PMLR, 2022.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoghu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented Transformer. arXio preprint arXio:2212.008136, 2022.
16

---

Published as a conference paper at COLM 2024
A Discussion: Selection Mechanism
Our selection mechanism is inspired by and related to concepts such as gating, hypemetworks, and data-dependence. It can also be viewed as related to "fast weights" (Schmidhuber, 1992; Ba et al., 2016), which connects classical RNNs with the mechanism of linear attention (Schlag et al., 2021). However, we believe that it is a distinct concept that is worth clarifying.
Gating. Gating originally referred to the gating mechanisms of RNNs such as the LSTM (Hochre- iter & Schmidhuber, 1997) and GRU (Chung et al, 2014), the gated equation in Theorem 1. This was interpreted as particular mechanism for controlling whether let an input into the hidden state of an RNN. In particular, this affects the propagation of signal through time and causes inputs to interact along the sequence length dimension.
However, the concept of gating has since been relaxed in popular usage to simply mean any multiplicative interaction (often with an activation function). For example, elementwise multiplicative components of neural network architectures (that do not interact along sequence length) are now commonly referred to as gated architectures (Hua et al., 2022; Mehta et al., 2023), despite a very different meaning than the original RNN sense. Thus we believe the original concept of RNN gating versus the popular usage of multiplicative gating actually have a very different semantic meaning.
Hypemnetorrss Hypernetworks refer to neural networks whose parameters are themselves generated by smaller neural networks. The original idea (Ha et al., 2017) used it in a narrow sense to define a large RNN whose recurrent parameters are generated by smaller RNN, and other variants have been around for a long time (Schmidhuber, 1992).
Data-dependence. Similar hypernetworks, data-dependence can refer any notion where some parameters of the model depend on the data (Poli et al., 2023).
Example: GLU Activation. To illustrate the issues with these concepts, consider a simple diagonal linear layer y Dx, where D is a diagonal weight parameter. Now suppose that D is itself generated from a linear transformation of x, with an optional nonlinearity: D=o(Wx). Since it is diagonal, the multiplication becomes an elementwise product: y (Wx))xx
This is a rather trivial transformation, yet it technically satisfies the common meanings of gating (since it has a multiplicative "branch"), hypernetworks (since the parameter D is generated by another layer), and data-dependent (since D depends on the data x). However, this in fact simply defines a GLU function, which is so simple that it is often considered just an activation function (Dauphin et al., 2017; Shazeer, 2020) instead of a meaningful layer.
Selection. Thus, while selection mechanisms could be considered special case of ideas such as architectural gating, hypemetworks, or data-dependence, so can an enormous range of other constructions eesiensallly anything with multiplication, including standard attention a mechanisms (Bahdanau et al., 2015; Vaswani et al., 2017) as well-and we find it uninformative to think of them as such.
Instead, we view it as most closely related to the gating mechanism of traditional RNNs, which a special case (Theorem 1) and also has a deeper history of connections to SSMs through variable finput-dependent discretization of A (Funahashi & Nakamura, 1993; Tallec & Ollivier, 2018; Gu et al., 2020a). We also eschew the term "gating" in favor of selection clarify the overloaded use of former. More narrowly, we use selection to refer to the mechanistic action of a model to select or ignore inputs and facilitate data interaction along the sequence length (Section 3.1). Beyond
selective SSMs and gated RNNs, other examples may include input-dependent convolutions (Yang et al., 2019; Lioutas & Guo, 2020; Kosma et al., 2023; Lutati et al., 2023) and even attention.
B Related Work
We overview several prior works related to our methods. We mention that some of the most closely related models include recurrent layers such as S4, S5, and quasi-RNNs; as well as end-to-end architectures such as H3, RetNet, and RWKV.
B.1 S4 Variants and Derivatives
We describe a brief overview of some structured SSMs from past work, particularly those that have relation to our method.
17

---

Published as a conference paper at COLM 2024
S4 (Gu et al., 2021; 2022a) introduced the first structured SSM, describing diagonal structure and diagonal plus low-rank (DPLR). It focused on efficient convolutional algorithms for DPLR SSMs due to a connection to continuous-time online memorization (HIPPO) (Gu et al., 2020a).
DSS (Gupta et al., 2022a) first discovered the empirical effectiveness of diagonal structured SSMs by approximating the HIPPO initialization. This was expanded on theoretically in S4D (Gu et 2022b).
S5 (Smith et al., 2023) independently discovered the diagonal SSM approximation, and is the first S4 model be computed recurrently with the parallel scan. However, this required lowering the effective state dimension, which they accomplished by switching the SSM dimensions from SISO (single-input single-output) to MIMO (mulltiinput mulli-outpuu) formulation. Our proposed S6 shares the scan, but differs by (i) keeping the SISO dimensions, which provides a larger effective recurrent state, (ii) using hardware-aware algorithm to overcome the computation issue, (iii) adding the selection mechanism.
Lu et al. (2023) applied S5 to meta-RL in order to handle resetting the SSM state between episode trajectories. Their mechanism can be viewed as a particular hard-coded instance a selection mechanism, where A is manually set to O, instead of our learnable mechanism that depends on the input. It would be interesting to apply selective SSMs generically to this setting and probe the model has learned to automatically reset its state on episode boundaries.
Mega (Ma et al., 2023) introduced a simplification of S4 to be real- instead of complex- valued, giv- ing it an interpretation of being an exponential moving average (EMA). They additionally make an interesting connection of the discretization step of SSMs to an EMA damping term. Contrary to findings in the original S4 papers, this was the first model to show that real-valued SSMs are em- pirically effective in certain settings or when combined with different architectural components.
Liquid S4 (Hasani et al., 2023) is also motivated by augmenting S4 with an input-dependent state transition. From this perspective it shares similarity to selection mechanisms, although in a limited form which is still computed convolutionally and close to LTI.
SGConv (Li et al., 2023), Hyena (Poli et al., 2023), LongConv (Fu et al., 2023), MultiresConv (Shi et al., 2023b), and Toeplitz Neural Network (Qin et al., 2023) all focus on the convolutional rep- resentation S4 and create global or long convolution kernels with different parameterizations. However, these methods cannot do fast autoregressive inference directly.
Notably, all of these methods, and all other structured SSMs that we are aware of, have been non-selective and usually strictly LTI (linear time invariant).
B.2 SSM Architectures
We use SSM architectures or state space neural networks (SSNN) to refer to deep neural network a architectures incorporating one of the previous SSMs as a black box layer.
GSS (Mehta et al., 2023) was the first gated neural network architecture incorporating SSMs. It is motivated by the gated attention unit (GAU) of Hua et al. (2022) and looks quite similar to our block, except with additional projections. Most importantly, its projection contracts the model dimension to reduce the state size of the SSM, while ours expands the model dimension in order to increase the state size, based on the motivation in Section 3.1.
Mega (Ma et al., 2023) combined the EMA simplification of S4 described above into a hybrid architecture using an efficient attention approximation.
H3 (Dao et al., 2023) is motivated by combining S4 with linear attention (Katharopoulas et al., 2020). It is the first to generalize this formulation of linear attention to more general recurrences, which is also the basis of later architectures.
Selective S4 (Wang et al., 2023) incorporates S4 as black box to generate binary mask which is multiplied on the input. While sharing the "selection" " name, we consider this an architectural modification that is closer to architectural gating than selection mechanism (Appendix A). For example, we hypothesize that it would not solve the Selective Copying task because simply masking out the irrelevant inputs does not affect the spacing between the relevant ones (indeed, the Selective Copying task can even be viewed as coming pre-masked if the noise tokens are embedded to 0).
RetNet (Sun al., 2023) is also based on Linear Attention and very similar to H3, but reduces the inner S4 layer to a special case where the state dimension is 11 Although not framed such, its recurrence can be viewed as special case of a linear SSM.
18

---

Published as a conference paper at COLM 2024
Its primary source of improvement is using linear attention with large head dimension, which can be viewed as another method to perform input-dependent state expansion. Using a larger head dimension in the context of linear attention variants was first done by H3, but not extensively used since this requires proportional amount of extra computation. RetNet avoids this with an alter- nate way to parallelize the computation with variant of standard multi-head attention instead of convolutions, made feasible by their particular special case of SSMs which acts as a simple EMA.
RWKV (Peng et al., 2023) is another recent RNN designed for language modeling, It is based on AFT (attention-free Transformer (Zhai et al., 2021)), another variant of linear attention. Its main "WKV" mechanism involves LTI recurrences and can be seen as the ratio of two SSMs.
We also highlight the gated attention unit (GAU) from Hua et al. (2022), which was motivated by combining the Transformer's 's MHA and MLP blocks together and was an inspiration for our architecture (Section 3.4) combining the H3 and MLP blocks.
B.3 Relationship to RNNs
RNNs and SSMs are broadly related, as they both involve the concepts of recurrence on latent state.
Several older RNNs such the strongly typed RNN (Balduzzi & Ghifary, 2016), quasi-RNN (QRNN) (Bradbury et al., 2016), and simple recurrent unit (SRU) (Lei et al., 2017; Lei, 2021) involve forms of gated RNNs without time-wise nonlinearities. Because of the connections of gating mechanisms and selection mechanisms, these can be viewed as cases of selective SSMs, and are thus more powerful in a sense than the family of LTI structured SSMs above. The main differences are:
They do not use state expansion (N= 1) or selective B,C parameters, both of which are important for performance (Appendix E.3).
They use a heuristic gating mechanism, which we generalize as consequence of the selection mechanism + discretization (Theorem 1). The connections to principled SSM theory provides better parameterizations and initializations.
Additionally, older RNNs famously suffered from efficiency issues and the vanishing gradients problem (Hochreiter, 1991; Hochreiter et al., 2001; Pascanu et al., 2013), both caused by their e sequential nature. The former could be solved for some of the above RNNs by leveraging the parallel scan (Martin & Cundy, 2018), but the latter was difficult without theory later developed for SSMs. For example, modern structured SSMs differ in more careful parameterization of the recurrent dynamics inspired by classical SSM theory (e.g. through discretization (Gu et al., 2021; 2023)), or direct analysis (Orvieto et al., 2023; Kaul, 2020; Gupta et al., 2022b)).
We also note that there is long line of work on orthogonal RNNs (Arjovsky et al., 2016; Henaff et al., 2016; Mhammedi et al., 2017; Vorontsov et al., 2017; Lezcano-Casado & Martinez-Rubit 2019) which are motivated by constraining the A transition matrix to be orthogonal or unitary, in order to control its eigenvalues and prevent the vanishing gradient problem. However, these had other limitations; we believe that these stem from the fact that orthogonal/ / RNNs are also LTT. For example, they are almost always evaluated on the Copying task which they can solve perfectly, but observed to struggle on the Selective Copying task (ling et al., 2019).
B.4 Linear Attention
The Linear Attention (LA) (Katharopoulos et al., 2020) framework is an important result populariz- ing kernel attention and showing how it relates to recurrent autoregressive models. Many variants have proposed alternative kernels and other modifications. Random Feature Attention (RFA) (Peng et all,,2022 2021) chooses the kemel feature map to approximate softmax attention (i.e. the exp feature map) using the random Fourier feature approximation of Gaussian kemness (Rahimi & Recht, 2007). Performer (Choromanski et al., 2021) finds an approximation to the exponential kernel involving only positive features, which also allows the softmax normalization term. Trrmsrormer (Qin et al., 2022a) showed that the LA denominator term can be unstable and proposed replacing it with LayerNorm. cosFormer (Qin et al., 2022b) augments RFA with cosine reweighting mechanism that incorporates positional information to emphasize locality. Linear Randomized Attention (Zheng et al., 2022) generalize RFA from the perspective of importance sampling, and generalize it to provide better estimates of the full softmax kemel (rather than just the exp- exp-transformed numerator).
Aside from kernel attention, many other variants of efficient attention exist; the survey Tay et al. (2022) offers an extensive categorization of many of these.
19

---

Published as a conference paper at COLM 2024
B.5 Long Context Models
Long context has become popular subject, and several recent models have claimed to scale to longer and longer sequences. However, these are often from a computational standpoint and have not been extensively validated. These include:
Recurrent Memory Transformer (Bulatov et al,, 2023), a lightweight wrapper around Transformer backbone. showed ability to generalize up to 1M sequences but only on synthetic memorization tasks; their main result is similar to our Induction Heads extrapolation experiment (Figure 3).
LongNet (Ding al., 2023), which claimed scale to 1B length but only evaluated on length <100K for actual tasks.
Hyena and HyenaDNA (Poli et al., 2023; Nguyen et al., 2023), which claimed to leverage up to 1M context. However, their experiments trained on proportionally more data at longer contexts, making it hard to conclude if quality improvements at 1M context are due to context length or due to more data and computation.
Sparse Transformer (Child et al., 2019) showed a proof-of-concept of using strided sparse attention Transformer to model audio waveforms of length 2 1048576, although did not discuss performance tradeoffs when controlling for computation and model size.
In contrast, we believe this work presents one of the first approaches to meaningfully demonstrate increasing performance with longer context.
C Mechanics of Selective SSMs
Prouf of Theurem 1. Consider a selective SSM (Algorithm 2) with N 1,4 1,B 1,sA Linear(x), = softplus. The corresponding continuous-time SSM (1)
h(f) = +x(+)
which is also called a leaky integrator.
The discretization step size is
A (Parameter +sx(x))
offtpis(PPrameter--Lierrrr)
softplus(Linear(x))
where we observe that the parameter can be viewed learnable bias and folded into the linear projection.
Now applying the zero-order hold (ZOH) discretization formulas:
A e
1 1+exp(Linear(1) = (((_iiear(xx))
1 11oo(Linar((x)))
B. = (AAA) (())e)ppAA-))-ABBBB 1). -(xxp(())) (O0(AA)))=111
=x((erar((x:))
Thus the final discrete recurrence (2a) is
8: = o ((inearr(xx))
h = -
as desired.
D Hardware-aware Algorithm For Selective SSMs
Without input-dependent selectivity,, SSMs can be efficiently implemented as a convolution (Gu et al., 2022a; Dao et al., 2023), which leverages the fast Fourier transform (FFT) as primitive. With selectivity, SSMs are no-longer equivalent to convolution, but we leverage the parallel associative scan. While SSM scans are theoretically efficient (O(BLDN) FLOPs, scaling linear in L), training foundation models with selective SSMs requires them to be efficient on modern hardware (GPUs) as well. We describe how we use kernel fusson and recomputation to make SSM scan fast and memorryeffficiet, We evaluate the speed of our scan implementation compared to convolution and attention in Appendix F.4, showing that it is up to 7x times faster than attention at sequence length 32K, and is as memory-efficient the best attention implementation (Flash Attention)...
20

---

Published as a conference paper at COLM 2024
Speed. On modern hardware accelerators (GPUs) most operations (except matrix multiply) are bounded by memory-bandwidit (Williams et al., 2009; Ivanov et al., 2021; Dao et al., 2022). This the case with our scan operation, and we use kernel fusion to reduce the amount of memory IOs, leading to significant speedup compared to a standard implementation.
The standard way to implement the scan algorithm in Section 3.2 is to prepare the scan input A,B of size (B,, in GPÚ HBM thigh-bandwidth memory, commonly referred to as GPU memory), call a parallel associative scan implementation to write the scan output of size (B,L,D,N) to GPU HBM, then multiply that scan output with C to produce an output of size (B, However, this requires the number of memory reads/writes /rrites on the order of O(BLDN). We can instead fuse the discretization step, the scan, and the multiplication with into one kernel:
1. We read in O(BLD+ bytes of memory (A,,,B,C) from slow HBM to fast SRAM.
2. We discretize to produce A,B of size (B,LD,N) in SRAM.
3. We perform a parallel associative scan, yielding intermediate states of size (B,L,D,N) in SRAM.
4. We multiply and sum with C, producing outputs of size (B,LD) and write it to HBM.
This way, we reduce IOs by a factor of O(N) (the state dimension), which in practice speeds up the operation by 20-40 times (Appendix E.4).
For sequence length L too long where we cannot fit the sequence in SRAM (which is much smaller than HBM), we split the sequences into chunks and perform the fused scan on each chunk. As long as we have the intermediate scan states, we can continue the scan with the next chunk.
Memory. We describe how we use the classical technique of recomputation to reduce the total amount of memory required to train selective SSM layers.
From the way we fuse the forward pass, we do not save the intermediate states of size (B,LDDN) to avoid memory blowup. However, these intermediate states are necessary for the backward pass to compute gradients. We instead recompute those intermediate states in the backward pass. Since the inputs A,A,B,C and output gradient read from HBM to SRAM are of size O(BLN + DN), and the input gradients are also of size O( DN), recomputation avoids the cost of reading O(BLND) elements from HBM. This means that recomputation of the SSM states in the backward pass speeds up the computation compared storing them and reading them from HBM.
Beyond optimizing for the memory requirement of just the scan operation, we also use recom- putation to optimize the memory requirement of the entire selective SSM block (input projection, convolution, activation, scan, output projection). In particular, we do not save intermediate activations that take lot of memory but are fast to recompute (e-g. output of activation function or short convolution). As result, the selective SSM layer has the same memory requirement as an optimized Transformer implementation with FlashAttention.. In particular, each attention layer (FlashAttention) stores around 12 bytes of activations per token, each MLP layer stores around 20 bytes of activations per token, for a total of 32 bytes ((assuming mixed -peccisinn training in FP16 or BF16)). Each selective SSM stores around 16 bytes of activations per token. Hence two layers of selective SSMs have around the same activation memory as an attention layer and an MLP layer.
E Full Experiments
E.1
Synthetic Tasks
Full experiment details for these tasks including task details and training protocol are in Appendix F.1.
E.1.1
Selective Copying
The Copying task is one of the most well-studied synthetic tasks for sequence modeling, originally designed to test the memorization abilities of recurrent models. As discussed in Section 3.1, LTT SSMs (linear recurrences and global convolutions) can easily solve this task by only keeping track of time instead of reasoning about the data; for example, by constructing a convolution kernel of exactly the right length (Figure 1). This was explicitly validated in earlier work on global convolutions (Romero et al., 2021). The Selective Copying task prevents this shortcut by randomizing the spacing between tokens. Note that this task has been introduced before the Denoising task (Jing et al., 2019).
21

---

Published as a conference paper at COLM 2024
Note that many previous works argue that adding architecture gating (multiplicatiive interactions) can endow models with dlata-dependence" and solve related tasks (Dao et al., 2023; Poli et al., 2023). However, we find this explanation insufficient intuitively because such gating does not interact along the sequence axis, and cannot affect the spacing between tokens. In particular architecture gating is not an instance of a selection mechanism (Appendix A).
Table 1 confirms that gated architectures such as H3 and Mamba only partially improve performance, while the selection mechanism (modifying S4 to S6) easily solves this task, particularly when combined with these more powerful architectures.
E.1.2 Induction Heads
Induction heads (Olsson et al., 2022) is a simple task from the mechanistic interpretability lens (Elhage et al., 2021) that is surprisingly predictive of the in-context learning ability of LLMs. requires models to perform associative recall and copy: for example, if the model has seen bigram such as "Harry Potter" in the sequence, then the next time "Harry" appears in the same sequence, the model should be able to predict "Potter" " by copying from history.
Dataset. We train a 2-layer model on the induction heads task at sequence length 256, with vocab size of 16, which is comparable to prior work on this task (Dao et al., 2023) but with longer sequences. We additionally investigate generalization and extrapolation abilities by evaluating range of sequence lengths from 2 = 64 up to 20 1048576 at test time.
Models. Following established work on induction heads, we use layer models, which allows
attention to mechanistically solve the induction heads task (Olsson et al., 2022). We test both multi-head attention (8 heads, with various positional encodings) and SSM variants. We use model dimension D of 64 for Mamba and 128 for the other models.
Results. Figure 3 shows that Mamba- or more precisely, its selective SSM layer-has the ability to solve the task perfectly because of its ability selectively remember the relevant token while ignoring everything else in between. It generalizes perfectly to million-length sequences, or 4000x longer than it saw during training, while no other method goes beyund 2x.
Out of positional encoding variants for attention models, xPos (which extrapolation) slightly better than the others; also note that all attention models were only tested up to sequence length 214 16384 due to memory limitations. Out of other SSMs, H3 and Hyena are similar, contrary to the findings in Poli et al. (2023).
designed for length
E.2 DNA Modeling
For pretraining, we largely follow standard causal language modeling (next token prediction) setup for the training and model details (see also Appendix F.2). For the dataset, we largely follow the setup of HyenaDNA (Nguyen et al., 2023), which uses the HG38 dataset for pretraining consisting of single human genome with about 4.5 billion tokens (DNA base pairs) in the training split.
E,2.1 Scaling: Model Size
In this experiment, we investigate the scaling properties of genomics foundation models with various model backbones (Figure 5 Left).
Training. To advantage the baselines, we train on a short sequence length of 1024; as shown in Appendix E.2.2, we expect results to favor Mamba even more at longer sequence lengths. We fix a global batch size of 1024, for a total of 200 &s 1M tokens per batch. Models were trained for 10K gradient steps for total of 10B tokens.
Results, Figure (Left) shows that Mamba's pretraining perplexity improves smoothly with model size, and that Mamba scales better than both HyenaDNA and Transformer++ For example, at the largest model size of zs 40M parameters, the curve shows that Mamba can match the Transformer++ and HNA models with roughly 3x to 4x fewer parameters.
E.2.2 Scaling: Context Length
In the next DNA experiment, we investigate the scaling properties of models with respect to sequence length. We only compare the HyenaDNA and Mamba models, quadratic attention
22

---

Published as a conference paper at COLM 2024
becomes prohibitively expensive at longer sequence lengths. We pretrain models sequence lengths 2 = 1024, 212 =096, 2 = 113844, 2 = 65536, 2 = 262144, 2 1048576. We fix model size of 6 layers by width 128 (about 1.3M-1.4M parameters). Models were trained for 20K gradient steps for total of x 330B tokens. The longer sequence lengths used sequence length warmup similar to (Nguyen et al., 2023).
Results. Figure 5 (Right) shows that Mamba able to make use of longer context even up to extremely long sequences of length 1M, and its pretraining perplexity improves as the context in- creases. On the other hand, the HyenaDNA model gets worse with sequence length. This is intuitive from the discussion in Section 3.5 on properties of the selection mechanism. In particular, LTI models cannot selectively ignore information; from a convolutional perspective, very long convolution ker- nel is aggregating all information across long sequence which may be very noisy. Note that while HyenaDNA claims to improve with longer context, their results do not control for computation time.
E.2.3 Synthetic Species Classification
We evaluate models on downstream task of classifying between 5 different species by randomly sampling a contiguous segment of their DNA. This task is adapted from HyenaDNA, which used the species /human, lemur, moo.se, wig.eiiiiiiipip h. We modify the task to be significantly more challenging by classifying between the five great apes species
[human, chwnaanze.mmaammgcccmem gor wrwgaaoggguannnoooo.ooo ,, which are known to share 99% of their DNA.
E.3
Model Ablations
We perform series of detailed ablations on components of our model, focusing on the setting of language modeling with size 350M models at Chinchilla token counts (same setting as Figure 4).
E.3.1
Architecture
Table 5 investigates the effects of the architecture (block) and its inner SSM layer (Figure 2). We find that
Among previous non-selective (LTI) SSMs, which are equivalent to global convolutions, performance is very similar.
Replacing the complex-valued S4 variant from previous work with a real-valued one does not affect performance much, suggesting that (at least for LM) real-valued SSMs may be a better choice when accounting for hardware efficiency.
Replacing any of these with selective SSM (56) significantly improves performance, validating the motivation of Section
The Mamba architecture performs similarly to the HB architecture (and seems slightly better when using selective layer).
We also investigate interleaving the Mamba block with other blocks such as MLP (a traditional architecture) MHA (a hybrid attention architecture) in Appendix F22.2
E.3.2 Selective SSM
Table 6 ablates the selective SSM layer by considering different combinations of selective A, B, and C parameters (Algorithm 2), showing that A is the most important parameter due its connection to RNN gating (Theorem 1).
Table 7 considers different initializations of the SSM, which have been shown to make a large difference in some data modalities and settings (Gu et al., 2022a;b). On language modeling, we find that simpler real-valued diagonal initializations (S4D-Real, row 3) instead of more standlrd complex-valued parameterizations (S4D-Lin, row 1) perform better. Random initializations also work well, consistent with findings from prior work (Mehta et al., 2023).
Table 8 and Table consider varying the dimension of the and (B,C) projections respectively. Changing them from static to selective provides the most benefit, while increasing the dimensions further generally improves performance modestly with a small increase in parameter count.
Of particular note is the dramatic improvement of the selective SSM when the state size N is increased, with over a 1.0 perplexity improvement for a cost of only 1% additional parameters. This validates our core motivation in Sections 3.1 and 3.3.
23

---

Published as a conference paper at COLM 2024
Table (Ablations: Architecture and SSM layer.) The Mamba block performs similarly to H3 while being simpler. In the inner layer, there little difference among different parameterizations of I.TI models, while selective SSMs (S6) provide large improvement. More specifically, the S4 (real) variant S4D-Real and
the (complex) variant S4D-Lin.
MODEL ARCH.
SSM LAYER
PPL
MODEL ARCH.
SSM LAYER
PPL
Hyena
H3
Hyena
10.24
Mamba Hyena
10.75
H3
H3
S4 (complex)
10.30
Mamba S4 (complex)
10.54
H3
S4 (real)
10.34
Mamba S4 (real)
10.56
H3
S6
8.95
Mamba Mamba S6
8.69
Table (Ablations:
Parameterization A.)
Table 6: (Ablations: Selective parameters.) A is the most important parameter (Theorem 1), but using multiple selective parameters together synergizes.
The more standard initializations based on
S4D-Lin (Gu et al., 2022b) perform worse than
S4D-Real or random initialization, when the SSM is selective.
SELECTIVE A SELECTIVE B SELECTIVE C PPL
X /
X
10.93
An INITIALIZATION FIELD
PPL
X
10.15
A, +ni
Complex 9.16
X
V
9.98
A. - f/2
Real
8.85
X
X
9.81
A (n+1)
Real 8.71
V
V
8.71
A. A)xexp(N(01))
Real 8.71
Table 8: (Ablations: Expressivity of A.) The selection mechanism of A constructs it with projection of the input. Projecting even to dim. provides large increase in performance, increasing it further provides further improvements at the cost of a modest increase in pa- rameters. State size fixed to N = 16.
Table (Ablations: SSM state dimension.) (Top) Constant and (Botfom) Selective B and C. Increasing the SSM state dimension N, which can be viewed as an expansion factor on the dimension of the recurrent state, can significantly improve performance for negligible cost in parameters/ Saaeettrrr///PPPP /LOPss but only when B and C are also selective. Size of A projection fixed 64.
STATE DIMENSION N PARAMS (M) PPL
1
367.1
9.88
SIZE OF A PROJ. PARAMS (M) PPL
2
367.4
9.86
4
368.0
9.82
358.9
9.12
8
369.1
9.82
1
359.1
8.97
16
371.5
9.81
2
359.3
8.97
4
359.7
8.91
1
367.1 1
9.73
8
360.5
8.83
2
367.4 4
9.40
16
362.1
8.84
4
368.0
9.09
32
365.2
8.80
8
369.1
8.84
64
371.5
8.71
16
371.5
8.71
E.4 Speed and Memory Benchmarks
We benchmark the speed of the SSM scan operation (state expansion N = 16), as well as the end-to-end inference throughput Mamba, in Figure 8. Our efficient SSM scan is faster than the best attention implementation that we know of (FlashAttention-22 (Dao, 2024)) beyond sequence length 2K, and up 20-40x faster than standard scan implementation in PyTorch. Mamba achieves 4-5x higher inference throughput than Transformer of similar size, since without the KV cache can use much higher batch sizes. For example, Mamba-6.9B (untrained) would have higher inference throughput than a 5x smaller Transformer-1 Details in Appendix F.5, which additionally includes benchmark of memory consumption.
24

---

Published as a conference paper at COLM 2024
Table 10: (Induction heads.) Models are trained on sequence length 2 236, and tested on various sequence lengths of 2 64 up to 2 1048576. / denotes perfect generalization accuracy, while X denotes out of memory.
MODEL
PARAMS
TEST ACCURACY (%) AT SEQUENCE LENGTH
26
2?
2
2
2
2
2
2
2
2
2
2
2
2
2
MHA-Abs
137K
99.6
100.0
58.6
26.6
18.8
9.8
10.9
7.8
X
X
X
X
X
X
MHA-RoPP
137K
V
100.0
83.6
31.3
184
8.6
9.0
5.5
X
X
X
X
X
MHA-xPos
137K
V
100.0
99.6
67.6
25.4
7.0
9.0
7.8
X
X
X
X
X
H3
153K
100.0
80.9
39.5
23.8
14.8
8.2
5.9
6.6
8.2
4.7
8.2
6.3
7.4
Hyena
69M"
97.7
V
100.0
/
441
12.5
6.6
5.1
7.0
5.9
6.6
5,9
6.3
9.8
Mamba 74K
100.0
/
/
/
/
Most the parameters are in learnable positional encodings.
Table 11: (Scaling Law Model Sizes.) Our model sizes and hyperparameters for scaling experiments. (Model dimension and number of heads applies only to Transformer models.)
PARAMS n_layers
d_model
n_heads / d_head TRAINING STEPS
LEARNING RATE BATCH SIZE
TOKENS
125M
12
768
12 / 64
4800
6e-4
0.5M tokens 2.5B
350M
24
1024
16 / 64
13500
3e-4
0.5M tokens 7B
760M
24
1536
16 / 96
29000
2.5e-4
0.5M tokens
15B
1.3
24
2048
32/ / 64
50000
2e-4
0.5M tokens
26B
F
Experimental Details and Additional Results
F11
Synthetic Tasks
Selective Copying, Our setting is on sequences of length 4096, with a vocab size of 16 possible tokens (including the white "noise" token from Figure 1) and requiring models to memorize 16 "data" tokens. We use layer models with a model dimension of D4
Models are trained for 400K steps at constant learning rate of 0.0001 with batch size of 64.
Induction Heads. Training consists of randomly generating data every step, with batch size of 8. We choose an "epoch" size of 8192 steps, and track the accuracy on fixed validation sets (also randomly generated) of each target sequence length. For the MHA-Abs and Mamba models, results are reported after the 25th epoch (8192 25 204800 steps). For the MHA-RoPE and MHA-xPos models, results are reported after the 50th epoch (8192x x 50 409600 steps). For the LTI H3 and Hyena models, results are reported after the 10th epoch (81920 steps) because they had converged by then and failed to improve further.
We use the Adam optimizer with no weight decay. All models are trained at constant learning rates 2e 4 and 1 3, and the better results reported for each model (2e 4 for all models except Mamba). The attention and Hyena models did not learn at LR 1e- 3, H3 learned at both LRs, but interestingly generalized better to shorter sequences at the smaller I.R of 2e - 4. Mamba learned both LRs, but extrapolated better the larger LR of 1e-3.
F22 Language Modeling
F22.1 Scaling Law Details
Scaling law experiments generally followed the GPT3 recipe. All models were trained on the Pile with the GPT2 tokenizer.
Model Sizes. Table 11 specifies the model sizes we use for scaling laws. This is taken directly from the GPT3 specifications (Brown et al., 2020), with very minor modifications, First, we changed the batch size of the 1.3B model from 1M tokens to 0.5M tokens, since we did not use enough parallelization to require the larger batch size. Second, we changed the number of training steps and total tokens to roughly match Chinchilla scaling laws (Hoffmann et al., 2022), which specify that training tokens should increase proportionally to model size.
Training Recipes All models used the AdamW optimizer with
25

---

Published as a conference paper at COLM 2024
gradient clip value 1.0
weight decay 0.1
no dropout
linear learning rate warmup with cosine decay
By default, the peak learning rate is the GPT3 specification.
We give several models an "improved recipe", inspired by changes adopted by popular large language models such as PaLM (Chowdhery et al., 2023) and LLaMa (Touvron et al., 2023). These include:
linear learning rate warmup with cosine decay to le- 5, with a peak value of 5x the GPT3 value
no linear bias terms
RMSNorm instead of LayerNorm
AdamW hyperparameter = (19,95) (the GPT3 value) instead of the PyTorch default of (99,999)
Architecture and Training Details. Our models are:
Transformer: The standard Transformer based on GPT3 (Table 11).
Transformer+4:: A Transformer with an improved architecture, namely rotary positional encod- ings (Su et al., 2021) and SwiGLU MLP (Shazeer, 2020), and the improved training recipe above. Hyena: Interleaving a Hyena block (the H3 block with S4 replaced by global convolution param- eterized by an MLP) with standard MLP blocks. The MLP blocks have expansion factor 2 instead of 4 and the number of layers is correspondingly increased by 1.5x to preserve parameter count.
H3++: The H3 architecture with a few modifications, including (i) using the same "thin" Hyena dimensions above (ii) the improved training recipe above (iii) linear attention head dimension of 8.
RWKV: The default RWKV model from Peng et al. (2023), including its modified MLP block. We also used as much of its specified training recipe as possible, such increasing the learning rates by 2x or 3x on certain parameters.
RetNet: The default RetNet model from Sun et al. (2023). We also gave it the improved training recipe above.
Mamba: The standard Mamba architecture, with the improved training recipe.
F.2.2
Additional Scaling Law Ablations
We perform additional ablations on the architecture using the same protocol as the 2k context length scaling laws in Figure 4 (Left).
Mamba Architecture: Interleaving Blocks. We test the effect of different architectural blocks combined with the Mamba block. We focus on the viewpoint that the Mamba block simply the standard SwiGLU block with an extra conv conv SSM path added. This leads to two natural ablations:
What the Mamba block is interleaved with a standard MLP block, instead of stacked homogenously? This can also be interpreted as taking Mamba and removing half of the SSMs.
What if the Mamba block is interleaved with MHA (multi-head attention) blocks? This can also be interpreted as taking a Transformer with SwiGLU MLPs (i.e. what we call Transformer++) and simply adding SSMs to the MLP blocks.
Figure 9 (Right) shows these variants compared to the original (homogenous) Mamba architecture. Interestingly, neither change matters too much. The Mamba-MLL architecture only slightly worse, and still better than all models except Transformer++ The Mamba-MHA architecture is only slightly better, which is somewhat surprising in light of the fact that many recent works have found that combining (LTT) SSMs with Attention can lead to substantial improvements (Dao et al., 2023; Fathullah et al., 2023; Saon et al., 2023; Zuo et al., 2022; Fathi et al., 2023).
26

---

Published as a conference paper at COLM 2024
(Seqquncce Length
Scaling (Bequence Length 2048)
Peplesty (bo stak)
Mambs-M4 Mareba 10HA
Perpleatty (to ecte
Hilee
1-1
7+337
12t
10
FLOPs
FLOPa
Figure (Scaling laws: extra ablations.) (Left) Instead of (Right) Instead
H3 Architecture: Training Recipes. Next we ablate differences between the Hyena and H3++ models, our weakest and strongest models outside of Transformer++ and Mamba, particularly to isolate the effect of training recipes.
Hyena: The Hyena block with its original architecture and GPT3 training recipe (same as Figure 4).
Hyena+: The same architecture but with the improved training recipe described above.
H3+: The same architecture Hyena+ but with the Hyena convolution kernel swapped out for S4D convolution kernel.
H3++: The same H3+, but with linear attention head dimension of 8. This increases computation inside the SSM recurrence but does increase parameters.
Our general convention is that "Model+" represents the base model with the improved training recipe, and "Model++" also allows for architectural changes.
Figure 9 (Right) shows that
A large improvement is achieved by the improved training recipe, which was used for many of the models in the main Figure (RetNet, H3++, Transformer++, Mamba).
The choice of the inner LTI SSM does not matter (e.g. Hyena vs. S4), consistent with findings throughout this paper.
The head dimension expansion improves performance, consistent with one of our main themes that expanded state dimension improves performance for SSMs (Section 3).
F2.3 Downstream Evaluation Details
This pretraining procedure is the same the scaling law protocol, but extended to 300B tokens and with the GPT-NeoX tokenizer (Black et al., 2022) instead of GPT2 tokenizer. For the 1.3B model, we use batch size of 1M tokens to be consistent with the GPT3 specifications. We report the perplexity on the Pile validation set, and for this metric only compare models trained on the same dataset and with the same tokenizer, in particular Pythia and RWKV.
For downstream evaluation, we use the LM evaluation harness from EleutherAl (Gao et al., 2021), as done by most work in this area. We evaluate the following tasks/ datasets that measure common sense reasoning:
LAMBADA (Paperno et al., 2016)
HellaSwag (Zellers et al., 2019)
PIQA (Bisk et al., 2020)
ARC-challenge (Clark et al., 2018)
ARC-easy: an easy subset of ARC-challenge
WinoGrande (Sakaguchi et al., 2021)
We report accuracy for LAMBADA, WinoGrande, PIQA, and ARC-easy, and accuracy normalized by sequence length for HellaSwag and ARC-challenge (since normalized accuracy is higher for almost all models for these task).
27

---

Published as a conference paper at COLM 2024
F3
DNA Modeling
F.3.1
Pretraining Details
We describe the dataset and training procedure of the HG38 pretraining task in more detail.
The dataset follows the splits from the prior Enformer work on genomics (Avsec et al., 2021); the training split contains a total of =34021 segments of length 2 = 131072 that cover the genome, for a total of approximately 4.5 billion tokens (DNA base pairs). These segments are pairs of (chromosome number, starting index, ending index), and can be extended necessary (e.g. to get longer segments).
We deviate from HyenaDNA when the training sequence length is not 217, HyenaDNA always takes a fixed sub-segment (eg. the beginning or middle of the prescribed segment), and thus for any training sequence length each epoch is fixed to 34021 samples and doesn't 't necessarily go through the whole genome. On the other hand, we use the entire training data:
When the context length I. less than (or equal to) 217, we divide up each segment into 2 I
non-overlapping sub-segments of length L, so that there are S total samples and 5 Sx2 4.5B tokens per epoch.
When the context length is greater than 217 we turn each segment into two samples, one that begins with the prescribed segment and one that ends with the prescribed segment. Thus each epoch has 2S items and 2SL tokens per epoch. For example, at sequence length 2 = =621444 there are 4x as many tokens as the default, and at sequence length 20 there are 16x as many tokens.
Other training details generally follow the same protocol as our language modeling experiments (Appendix F.2). For example, we use the AdamW with (f1,82) = (0.9,0.95), no dropout, weight decay 0.1. We use cosine learning rate scheduler with linear warmup for 10% of total steps.
F33.2
Scaling: Model Size Details
Models.
The models we consider are:
Transformer++: Transformer with improved architecture, notably the usage of RoPE positional encodings (Su et al., 2021). Informally, we found these to be noticeably better than vanilla positional encodings from (Vaswani et al,, 2017).
HyenaDNA: the Hyena model from Poli et al. (2023); Nguyen et (2023), which is roughly Transformer with the MHA block replaced by an H3 block using global convolution parameterized by an MLP.
Mamba: the standard Mamba architecture.
Model Sizes.
We use the following model sizes.
BLOCKS MODEL DIMENSION
4
5
6
7
8
10
12
64
96
128
192
256
384
512
PARAMS (APPROX.) 250K
700K
1.4M
3.5M
7.0M
19.3M
40.7M
Note that the number of blocks for Mamba is doubled, because one Transformer "layer" includes both the MHA and MLP blocks (and similarly for Hyena), which requires two Mamba blocks to match parameters (Section 3.4).
Training. For each model (Transformer++, HyenaDNA, Mamba), swept the learning rate across (le-3,2e 3,2e-3,4e 3,8e-3), 3). The optimal Transformer and HyenaDNA learning rates were 2e-3 across all sizes. The optimal Mamba learning rate was 8;-3; note that Mamba performed better than baselines with matched learning rates (2e-3), but was more stable and improved even more at higher learning rates. (Furthermore, as this LR is on the upper range of the sweep, it possible that our results are still suboptimall)
Note that, in contrast standard LM scaling laws (Table 11), our LR held constant across model sizes for simplicity. The optimal LR should go down for larger models, but we didn't find noticeable effect at the small model sizes (at most few million parameters) considered.
28

---

Published as a conference paper at COLM 2024
F3.3 Scaling: Context Length Details
We use a total batch size of 2 z 16M tokens per training step, for every sequence length (e.g. at length 2 there are 16 segments per batch and at length 210 there are 16384 segments per batch). This is large batch size relative to the model size by usual LM standards, but note that a batch size of 2 is the minimum possible on machine with 8 GPUs and sequence length of 2-0,, and that HyenaDNA used much larger batches of 2
The learning rate used was 0.008 for Mamba and 0.001 for HyenaDNA; we initially attempted to use the same learning rate of 0.002 from the previous section for HyenaDNA, but found that it was unstable at the longest context length.
Sequence Length Warmup. Following (Nguyen et al., 2023), we use sequence length warmup (SLW) during pretraining. We choose a simple schedule of 2 epochs each power-of-two sequence length starting from 2 1024. (Note that because of how data is curated, at the longest sequence lengths more steps and tokens are spent proportionally. In particular, each stage up to length 2 processes the same number of tokens, but 4x as many tokens are processed at length 218, 8x as many at length 2 and 16x as many at length 20
Unlike HyenaDNA, we always control for the number of tokens per gradient update, so the batch size successively halved as the sequence lengths are doubled in each stage.
Remark F.1. We also note that the schedule was not tuned, and we never experimented with turning off sequence length warmup for these pretraining experiments. We later found that SIW did not help noticeably for audio pretraining at similar lengths (Section 4.4), and it is possible that it is not necessary for DNA pretraining either.
E3.4 Species (Great Apes) Classification
Models are causal and therefore only the last element (across the sequence length) of the model's output used for the classification head. Note that we control for the total number of elements in the loss function per gradient step. The pretraining objective includes all positions across the sequence length, so that batch_size x sespuuenceeeegthhh is held constant; in other words, the batch size decreases as the sequence length increases. However, for a classification task, since only the last position enters the loss, the batch size itself is held constant. Note that this also means that fine-tuning models with longer sequence lengths is more computationally expensive.
Training consists of epochs, each of which has 1024 gradient steps. Fach gradient step uses batch size 64, which are all independently randomly drawn by uniformly picking species, uniformly picking a chromosome, and then uniformly picking contiguous segment of DNA.
Following (Nguyen et al., 2023), models with maximum context length greater than 2 16384 use sequence length warmup with 1 epoch at length 2 = 16384, epoch at length 2 - 32768,, 1 epoch at length 26 - -65336, and so on up to the maximum sequence length. For example, the model with 220 1048576 context undergoes 6 epochs of sequence length warmup before 4 more epochs at its maximum sequence length.
The learning rate for all Hyena models 4e while the learning rate for all Mamba models is 1e-4. These were found by performing learning rate sweeps for each model among 5,2e- 5,4e- 4,2e 4) for the smaller sequence lengths (210 222 214 22)), and these values were consis- tently found be the best for each model. An abridged learning rate sweep was done at length 18 which agreed with these values, and single run at length 220 was performed (as described above, the computational cost of these experiments is proportional to the sequence length). The learning rate followed cosine decay schedule with warmup with 5 epochs of linear warmup to the max- imum learning rate, and 5 epochs of cosine decay down to 1e-6. The unusually long learning rate warmup schedule was chosen because the sequence length warmup was also long (eg comprising 6 out of 10 epochs for the model with context length 200), we did not experiment with this choice.
Results for the Species classification task are in Table 12.
29

---

Published as a conference paper at COLM 2024
Table 12: (Great Apes DNA Classification.) Accuracy after fine-tuning on sequences of length 2 up to 2 1048576 using pretrained models of the same context length. Random guessing is 20%.
1024
MODEL
PARAMS ACCURACY (%) AT SEQUENCE LENGTH
210
212
2
2
2
2
HNenaDNA
1.4M
28.04
28.43
41.17
42.22
31.10
54.87
Mamba
14M
31.47
27.50
27.66
40.72
42.41
71.67
Mamba
7M
30.00
29.01
31.48
43.73
56.60
81.31
Table 13: YoulubeMix length scaling sequence lengths and batch sizes.
SEQUENCE LENGTH BATCH SIZE
TOKENS / BATCH
168x2008=958444 1
958464
234x2048 = 479232 2
958464
x17x 2236616 4
958464
19x2048 = -120832 8
966656
x 2048 -61440 16
983040
15x 2048 =30720
32
983040
8x 2048 = 16384
64
1048576
4x 2048 8192
128
1048576
F4
Audio Details
F4.1
YouTubeMix Audio Pretraining
Model. We use a model with 3 blocks per stage (3x5 5= 15 total Mamba blocks), pooling factor p= 16, and outer dimension D 64, for about 3.5M parameters.
The data is mu-law encoded 8 bits, so the model modeling discrete tokens with a vocab size of 256.
Dataset.
The dataset consists of clips of up to minute long, or length 960000, which is subsampled and divided into segments of any desired sequence length. Since the architecture involves two stages of pooling by a factor of 16, and we want the resulting sequence length to be a multiple of 8 for hardware efficiency, the longest possible sequence is 468 2048 1 958464. The rest of our sequence lengths are defined by successively halving this and rounding up to the nearest multiple of 2048.
Table 13 lists the specifications used in Figure 7. Beyond the varying batch sizes, the number of valid segments in the training set varied between different sequence lengths (e.g, the number of training steps per epoch was not constant for different points in the graph), which may have contributed to kinks in the scaling curves.
Training. Models were trained for 200K training steps with maximum learning rate of 0.002, 20K (10%) warmup steps, and weight decay 0.1 (similar to our general pretraining recipe across domains).
Additional Ablations: SSM Parameterizations. We investigate SSM parameterizations on forg-form&a audio waveform pretraining in the setting of Figure 7. The setting is modified slightly to use larger models (8 layers and D = 64 for 6M params, the SaShiMi default), shorter sequences (2 =248 to 2 2621 44 instead of 213 to 22t), lower LR (0.001 from 0.002), and shorter training cycles (100K instead of 200K steps).
Figure 10 shows that the change from S4 + S6 (i.e. the selection mechanism) is not always beneficial. On long-form audio waveforms, in fact significantly hampers performance, which may be intuitive from the point of view that audio is uniformly sampled and very smooth, and therefore benefits from continuous linear time-invariant (LTI) methods. After ablating away the selection mechanism, note that the resulting model is the S4 layer inside the Mamba block. To disambiguate, call this Mamba-SA S4 as opposed the default Mamba architecture Mamba-S6.
30

---

Published as a conference paper at COLM 2024
Audo Wereforms
Maria -aoeppix
1.45
noomplic
Bits Byee
- - (Wantax81)
Bltss Per Byee
1,40
vellccre A
(Mambe 54)
1.30
1.30
1.29
Sequence
Langth
Figure 10: (Audio Pretraining (YouTubeMix) Ablations.) As a unifornly-sappled "continuous" signal modality, audio waveforms actually benefit from LTI models which have matching inductive bias. (Left) Homogenous models (all blocks have the same parameterization) (Right) Only the center U-Net blocks are ablated; the outer blocks are Mamba-S4. Purple line is same as figure on left.
However, on the right side, we keep the outer layers of the U-Net Mamba-S4 and ablate only the inner layers. The performance differences shrink dramatically; this reinforces the hypothesis that layers closer the naw audio signal should be LTI, but once they are "tokenized" and compressed by the outer layers, the inner layers no longer need to be LTI. In this setting however, the real-valued SSM still underperforms the complex-valued one.
F.4.2
SC09 Speech Generation
Autoregressive training largely followed the autoregressive language modeling protocol, such as
Weight decay 0.1
Learning rate warmup for 10% of total steps
AdamW optimizer with B = (0.9,0.95)
Gradient clip value 0.1
We used learning rate of 0.002 and 200000 training steps at a batch size of 16.
The large Mamba model in Table has 15 layers per stage with an outer dimension of D % and pooling factor 4. We note that this dataset is small (training went through 100 epochs) and for this large model, there was significant overfitting of the BPB NLL. However, automated metrics of generated samples continually improving throughout training.
models in the architecture ablations in Table 4 all have 8 layers per stage with an outer dimension of 64 and pooling factor 4. The S4+MLP block has roughly 2D2++DD parameters (expansion fac- tor 2 in the MLP). The Transformer block has 4D2 +2D parameters (expansion factor in the MLP). The Mamba block has the usual x 6D2 parameters, All models have roughly 6M total parameters.
E5 Efficiency Benchmark
Scan Operation. We compare the core operation of selective SSMs, which is the parallel scan (Section 3.3), against convolution and attention, measured on an A100 80GB PCle GPU. Note that these do not include the cost of other operations outside of this core operation, such as computing the convolutional kernel in global-convolution models, or computing the QKV projections in attention.
As a baseline, we implement a standard parallel scan in PyTorch with no kernel fusion. This requires materializing the parameters A,B,C in HBM.
Our scan implementation fuses the discretization step and the parallel scan, avoiding the cost of materializing all the large parameters in HBM.
For convolution, we use the standard implementation in PyTorch, which separately performs FFTs on the inputs and the filters, multiply them in frequency domain, then performs an inverse FFT to obtain the result. The theoretical complexity is (Llog(()) for sequence length I..
For attention, we compare against the fastest implementation that we are aware of (FlashAttention- 2 (Dao, 2024)), with causal mask. Note that HlashAttention-2 with causal mask is about 1.7x faster than without causal mask, since approximately only half of the attention entries are computed.
We use batch size of 1 and increase the sequence length from 2 5122, 2 ~ 1K, 2 ~ 2K, up to 219 500K (some of the baselines out of memory before reaching 500K). We model
31

---

Published as a conference paper at COLM 2024
Table 14: (Memory benchmark.) Mamba's memory footprint comparable the most optimized Transformer. Results for 125M models.
Batch size Transformer (w/ FlashAttention-2) Mamba
1
4.6GB
4.8GB
2
5.2GB
5.8GB
4
6.9GB
7.3GB
8
11.5GB
12.3GB
16
20.7GB
23.1GB
32
34.5GB
38.2GB
dimension of = 1024 and state dimension N = 16. We measure with BF16 inputs, which is the data type most commonly used for large scale training.
End-to-end Inference. We measure the inference throughput of a Mamba 1.4B model and an untrained Mamba 6.9B model, against standard Transformer (GPT3 architecture) at 1.3B and 6.7B size. We use the standard Transformer implementation in the Huggingface transformers library.
We set the prompt length to be 2048 and the generation length to be 128. We vary the batch size from 1, 2, 4, 16, 32, 64, to 128, and measure time time taken to generate 128 tokens. We then calculate the throughput (tokens/s s) as batch size 128/ time taken. We repeat the measurements 3 times and take the average. Measurements done
A100 80GB PCIe GPU.
Memory Benchmark. The memory usage simply scales proportionally the size of the activation tensors, as with most deep sequence models. We report measurements of the training memory requirements of 125M models on A100 80GB GPU. Each batch consists of sequences of length 2048. We compare to the most memory-efficient Transformer implementation we are aware of (with kernel fusion from torch. compile le and with Attentioon))))2)) Table 14 shows that Mamba's memory requirement is comparable to a similar-sized Transformer with an extremely optimized implementation, and expect further improvement in Mamba's memory footprint in the future.
32


Reviews:
[{"review_id": "FBlG83q0ek", "reviewer": "Reviewer_YEh1", "summary": "The paper proposes an improvement over state-space models that solve issues resulting from static nature of SSMs. The main idea is that the parameters of SSMs (e.g., CNN filters) are not input-dependent. The proposed method gives the model the ability to solve tasks that require \"context-aware reasoning\".\n\nWhile this change sacrifices the efficiency of static SSMs, the paper also proposes a hardware kernel to speed up the models. The model demonstrates strong results on various modalities including text, DNA, audio etc. \n\nOverall, the paper presents a very interesting model that has clear motivations, with solid executions. However, the experiments results, especially the language part look insufficient to show its superiority over transformers.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the feedback and questions!\n\nWe would like to note that we view this work as re-invigorating interest into alternative architectures, and opening new directions in this space. The reviewer raises many interesting questions about other tasks/applications and potential efficiency improvements; while we could not address all of them in this first work, we are optimistic that much more progress and empirical validation will continue to be made.\n\n> Long-context tasks\n\nTo our knowledge there are not standard and high-quality benchmarks that are widely agreed to measure long context in text. Thus our goal was to focus on more standard LM pretraining objectives, as well as other modalities (e.g. DNA) where long context matters.\nWe note that follow-ups have further investigated long-context tasks, such as needle-in-a-haystack [LongMamba] and question-answering [Jamba], where Mamba (or variants) perform well.\n\n> Training efficiency\n\nWe note that Mamba is much faster than Transformers in wall-clock time for longer sequences.\nHowever, as the reviewer observes, on a *FLOP-for-FLOP* basis, it is still less hardware-efficient than Transformers because Mamba does not utilize as many matrix multiplications, which GPUs are extremely specialized for.\n\nThis is in fact an area of active research with very positive progress. For example, several follow-up works have been able to develop closely related models that leverage matrix multiplications and are much faster than Mamba [GLA, HGRN2, Mamba-2], comparable to Transformers in wall-clock time even for shorter sequences.\n\n> Fixed state size.\n\nIndeed, the tension between state size and recall ability is a fundamental tension of sequence models. We expect SSMs and Transformers to have different tradeoffs in inductive bias and capabilities. For example, SSMs naturally have more trouble with rote memorization [1], but have better ability to learn other types of formal languages [2]. In fact, we view SSMs and attention as complementary, and follow-up work has found that hybrid architectures are broadly very strong (e.g. Jamba, Zamba, Griffin).\n\n[1] https://arxiv.org/abs/2402.01032\n[2] https://arxiv.org/abs/2405.17394\n\n> Static A\n\nThe discrete A dynamics $\\bar{A}$ (which govern the SSM recurrence in eq. (2a)) depend on both $A$ and $\\Delta$. Since the latter is context-aware, so is $\\bar{A}$, which is the true recurrence. In early explorations, we found that additionally making $A$ input-dependent did not help performance.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The method is a solid improvement over S4. The results on synthetic tasks clearly highlight the drawbacks of vanilla SSMs and validated the improvements. \n* The efficient implementation proposed along with the model is quite interesting and seems useful to speed up other recurrent-style models. \n* The experimental results cover a range of modalities, demonstrating the models' potential in diverse applications.", "reasons_to_reject": "* Lack of results on long-context text tasks to justify that the models can actually replace transformers. \n* It is unclear whether the comparison with Transformers is totally fair, e.g., training cost in terms of wall clock time even though the Flops might be smaller, if the learning rate schedule has been properly tuned for transformers.\n* While the model has input-dependent behaviors, it still has a fixed size memory/hidden state of all previous context. This suggests that it might not work well in LLMs scenarios where the users can have different queries over the same context after the context has been encoded.", "questions_to_authors": "How good can the model be parallelized and utilize machine compute?\n\nCan you explain why only B, C needed to be context-aware and the static A can still be used?\n\nThe writing of the paper requires a good understanding of the precious S4 models, and I think it is worthwhile to add more context to the main paper instead of leaving them in the appendix. If space is an issue, I suggest the authors move the experiments of other modalities to the appendix given their low relevance to this conference.", "extra_scores": {"ethics_concerns_details": "N/A"}}, {"review_id": "c0GHou5rcy", "reviewer": "Reviewer_hthE", "summary": "The authors integrate selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5× higher throughput than Transformers) and linear scaling in sequence length, and its\nperformance improves on real data up to million-length sequences. Mamba achieves state-of-the-art performance across\nseveral modalities such as language, audio, and genomics. On language modeling, Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.", "questions": "", "limitations": "", "rating": 9, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very glad that the reviewer enjoyed the paper and model. Thanks for the encouragement!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Mamba not only solves synthetic tasks such as copying and induction heads easily but also can extrapolate solutions indefinitely long (>1M tokens).\n2. Mamba out-performs prior state-of-the-art models such as SaShiMi, Hyena, and Transformers on modeling audio waveforms and DNA sequences, both in pretraining quality and downstream metrics\n3. Mamba is the first linear-time sequence model that truly achieves Transformer quality performance, both in pretraining perplexity and downstream evaluations.\n4. Mamba is a hardware-aware algorithm. The resulting implementation is faster than previous methods.\n\nPersonally tested Mamba and it's amazing as the authors claimed.", "reasons_to_reject": "Don't find a good reason to reject.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XSB8vvovkv", "reviewer": "Reviewer_F6hZ", "summary": "This paper proposes Mamba, a linear-time sequence model with an intra-layer combination of Selective S4D, Short Convolution and Gated Linear Unit. The paper introduces input-dependent gating to SSM and provides a hardware-aware efficient implementation for the selective parallel scan operator. Mamba shows strong performance in language, DNA and audio modeling compared to models based on attention/RNN/convolution.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed feedback and insightful questions.\n\n1. Ablation Study: The short conv is a simplification of H3’s “shift SSM”. RWKV has “token shift” (equivalent conv with width=2). The short conv does improve model quality (at 350M with 7B tokens, ppl 8.71 with conv and 8.97 without conv). The primary improvement comes from selectivity (Table 6).\nExtensive ablation experiments are included in Appendix E.3 (in Table 5, 6, 7, 8, 9).\n\n2. Baselines:  We compared to Transformer++ in our scaling law study, showing that Mamba can match Transformer++ on perplexity. For table 2 (models up to 3B trained for 300B tokens), we compared to public models trained with similar numbers of tokens. Due to compute budget, we did not train our own Transformers. The fact that an SSM ( linear in seqlen) can match the strongest Transformer (quadratic in seqlen) and outperform other Transformer variants, which has gone through 7 years of intense tuning from the whole research community, is very encouraging to us.\n\n3. Training Efficiency:  Mamba has the same number of layers as Transformers (e.g. Mamba 2.7B has 64 S6 layers, Pythia 2.8B has 32 MLP and 32 attn layers, same total params). Mamba’s training cost is comparable to that of Pythia (Pythia 2.8B takes 14240 A100 hours, Mamba-2.7B takes 14370).\n\n4. Relationship to other input selection mechanisms:  We view this as analogous to “hard attention” (e.g. in Mnih et al., 2024 \"Recurrent models of visual attention\"), and “soft attention” (e.g. Bahdanau et al., 2015 “Neural machine translation …”), nowadays just called “attention” (e.g. by Vaswani et al., 2017). \n\n5. Length extrapolation: We will make clear that our claims are with respect to specific empirical results in (1) synthetic tasks (2) audio (3) genomics. Follow-up works have further investigated long-context tasks, such as needle-in-a-haystack (https://github.com/jzhang38/LongMamba) and more realistic NLP tasks such as question-answering (Jamba), where Mamba (or variants) perform well.\n\n6. Settings of the scaling law experiments: \nFull experiment settings are in Appendix F. For Transformer FLOPs: 6 * nparams * ntokens + attention FLOPs (12 * nlayers * hdim * ctx_len * ntokens). For Mamba: 6 * nparams * ntokens + scan FLOPs (27 * nlayers * hdim * state_dim * ntokens).\nThe Chinchilla (20 tokens per params) protocol was originally tuned for Transformers, and we do not change this (1) due to compute budget (2) to be generous to the Transformer baselines.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The strong performance of the model on both synthetic and the real world tasks.\n2. The proposed Selective State Space Models and its efficient implementation are significant contributions to the community of SSM and linear recurrence.", "reasons_to_reject": "1. Lack of Ablation Study: Besides the proposed Selective S4D module, the Mamba layer also includes the Short Convolution module which can also be added to other linear recurrent models. It is not clear how Short Convolution contributes to the performance of Mamba.\n2. Weak Baselines: The table 2 does not include the state-of-the-art transformer++ architecture that is used by Llama models, and the included baselines either have suboptimal architectural configurations or bad optimization/initialization recipes.\n3. Lack of Training Efficiecy Measurement: The strong performance of Mamba can be due to the fact that it has far more layers than other models. The authors should show the training cost of Mamba when compared with other language models in Table 2.\n4. The relationship to other input selection mechnism: While the authors claims that they \"use selection to refer to the mechanistic **action** of a model to select or ignore inputs\", Mamba does not make any discrete actions to select the input, but rather uses a soft gating techinque similar to Gated Recurrent Units. If the authors still want to use the terminology of selection, they should discuss its relationship to the works [1, 2, 3, 4] that have real discrete selections of inputs.\n5. The length extrapolation experiments are only conducted on the synthetic setting. The synthetic induction head task is too simple and the authors make a strong claim of perfect length extrapolation. However, this doesn't even translate to the language modeling setting where Mamba is known to be not extrapolable as tested in [5] and [6]. I would suggets authors to (i) conduct experiments on a more challenging synthetic setting such as Multi-Query Associative Recall (MQAR), and (ii) avoid claiming the infinite extrapolation ability of Mamba.\n6. The settings of the scaling law experiments are unclear and somewhat misleading. See question 6 below for more details. Also, in Figure 4, while increasing the sequence length from 2K to 8K, Mamba seems to have a worse scaling trend compared with Transformer++ when the training FLOPs is larger than $10^{20}$.\n\n---\n[1] Efficient Content-Based Sparse Attention with Routing Transformers (TACL 2020)\n\n[2] Reformer: The Efficient Transformer (ICLR 2020)\n\n[3] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (JMLR 2022)\n\n[4] Sparse Modular Activation for Efficient Sequence Modeling (NeurIPS 2023)\n\n[5] https://github.com/jzhang38/LongMamba\n\n[6] Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024)", "questions_to_authors": "1. How is the training wall time of Mamba compared to other models of similar sizes?\n2. How is the performance of Mamba compared to the Llama architecture on downstream NLP tasks?\n3. Is the short convolution module critical to Mamba's performance? What if we remove it?\n4. It seems that in the Mamba official implementation, the $\\mathbf{B}$ parameter in S6 is not using discretization, but rather a simple outer product. How does this discrepency affect the performance?  Does this indicate discretization is not neccessary for gated linear recurrent models? What about removing discretization for A? How is the discretization important generally?\n5. Why is Mamba not length extrapolable in real-world language modeling?\n6. How is the FLOPs calcuated in the scaling law experiments? Why using the Chinchilla protocol which is built for Transformers instead of fitting a new curve for Mamba?", "extra_scores": {}}, {"review_id": "c0dsSRH5EG", "reviewer": "Reviewer_pB6t", "summary": "Previous state-space models (SSMs) aimed to address the inefficiencies of Transformer architectures over long sequences, but they generally underperformed compared to Transformers. This paper introduces a new SSM architecture, Mamba, which has efficient long-context language modeling with linear complexity. Compared to S4, Mamba incorporates a selection mechanism that allows the model to selectively memorize information based on previous inputs. The architecture also combines Mamba layers with MLP blocks to further improve computational efficiency. The implementation uses a hardware-aware design that efficiently utilizes SRAM for discretization and recurrence. \n\nIn synthetic tasks such as selective copying and induction heads, Mamba demonstrated significant improvements over previous SSMs, including H3 and S4 models. These enhancements show Mamba's capability to handle long context effectively. Further experimentation with models scaling up to 2.8 billion parameters confirmed Mamba's robustness on language modeling, showing that it achieves comparable performance to Transformer models with the same amount of training data. The authors have conducted experiments extends beyond language modeling with experiments on audio and genomics data, which illustrates its great performance in processing long sequences across different modalities.", "questions": "", "limitations": "", "rating": 10, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the encouragement!\n\n> State compression and memory capacity\n\nAbsolutely – we believe that the tradeoff between state size and model capacity is central to sequence models, and there are many exciting research directions leading from here.\n* The Based paper [1] found that, when plotting memorization capability against effective state size, that various finite-state models including SSMs actually have comparable performance to attention (as a function of state size). Of course, pure softmax attention can get larger effective state sizes.\n* Recent SSM follow-ups have found ways to get much larger state sizes (GLA [2], HGRN2 [3], Mamba-2 [4], RWKV-6 [5], xLSTM [6]), which significantly improves their performance on synthetic memorization tasks.\n* We think that attempting to theoretically characterize the memory capacity of models in terms of their effective state size (e.g. proving a lower bound) would be a very interesting theoretical question on sequence models.\n* Another interesting direction for improving models is whether a model’s state size can grow with the context length. For example, the linear growth (i.e. KV cache) of a pure Transformer might be wasteful, whereas the constant size state of a pure SSM might be too restrictive; can a model’s state grow slowly through the sequence length instead?\n\n[1] Simran Arora et al. “Zoology: Measuring and Improving Recall in Efficient Language Models” ICML 2024.\n\n[2] Songlin Yang et al. “Gated Linear Attention Transformers with Hardware-Efficient Training” ICML 2024.\n\n[3] Zhen Qin et al. “HGRN2: Gated Linear RNNs with State Expansion” https://arxiv.org/abs2404.07904\n\n[4] Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality” ICML 2024.\n\n[5] Bo Peng et al. “Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence” https://arxiv.org/abs/2404.05892\n\n[6] Maximilian Beck et al. “xLSTM: Extended Long Short-Term Memory” https://arxiv.org/abs/2405.04517\n\n\n> Initial states and length generalization\n\nWe are very excited about passing states through the model so that forms of truncated back-propagation can be used to potentially learn on indefinitely long sequences. This functionality and other large improvements will be open-sourced to continue contributing to the community!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* Mamba has solved the implementation difficulty of previous SSMs with fused kernel and hardware-aware algorithm. \n\n* The authors have conducted extensive experiments on Mamba with different settings.\n\n* Mamba has greatly contributed to the development of efficient LLMs and SSMs. It has the potential to become the successor of Transformers.", "reasons_to_reject": "There is no reason to reject this paper.", "questions_to_authors": "* Compression is a key concept in SSMs and recurrent models, as history needs to be stored in fixed-size states, but it would limit the memory capacity. For Transformers, KV cache can grow infinitely, so theoretically they have unlimited memory capacity. It is interesting to see what is the upper-bound of Mamba's memory capacity in terms of memorizing history information.\n\n*  The current implementation doesn't support passing in initial state, and getting their gradients during back-propagation. This may help to solve the length generalization problem of Mamba, and potentially can be combined with gradient check-pointing in the sequence dimension to train the model with very long sequence length.", "extra_scores": {}}]

Author rebuttal:

Meta-review: Mamba addresses key architectural and implementational problems in S4, creating a potentially viable alternative to the transformer family of architectures. While the jury may still be out on what is the optimal architecture for modeling long sequences, this is a significant advancement in sequence modeling with sub-quadratic complexity.

As reviewers YEh1 and F6hZ suggest, additional evaluations, ablations, and comparisons would have cemented the claims even more; I also acknowledge that this is very difficult to do in an academic setting. The paper presents sufficient results in my view, and as the authors and additional reviewers have noted, there is a growing amount of external evidence and adoption to support the paper's claims.
Decision: Accept