Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 22:38:43

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.

Main paper: 
 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 a wide variety of domains such as 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 to its ability to 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 a 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, 2022; 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 an 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 a 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 IO 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-based SSMs) and on modern hardware (up to 3× faster on A100 GPUs).Architecture. 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 a single block, leading to a simple and homogenous architecture design (Mamba) incorporating selective state spaces.Selective SSMs, and by extension the Mamba architecture, are fully recurrent models with key properties 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 genomics. (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 a 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 as 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-length sequences.• Language Modeling. Mamba is the first linear-time sequence model that truly achieves 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 a large range of baselines, including very strong modern Transformer training recipes based on LLaMa (Touvron et al., 2023) . Our 1.4B Mamba language model has 5× inference throughput compared to Transformers of similar size, and its quality matches that of Transformers twice its size (e.g. 5 points higher avg. on common sense reasoning compared to Pythia-1.4B and matching Pythia-2.8B). 
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 a particular continuous system (1) that maps a 1-dimensional function or sequence x(t) ∈ R → y(t) ∈ R through an implicit latent state h(t) ∈ R N . Concretely, S4 models are defined with four parameters (∆,A,B,C), which define a sequence-to-sequence transformation in two stages.h ′ (t) = Ah(t)+Bx(t) y(t) = Ch(t) (1) h k = Ah k-1 +Bx k (2a) y k = Ch k (2b) K = (CB,CAB,...,CA k B,...) (3a) y = x * K (3b)Discretization The first stage transforms the "continuous parameters" (∆, A, B) to "discrete parameters" (A,B) through fixed formulas A = f A (∆,A) and B = f B (∆,A,B), where the pair(f A , f B ) is called a discretization rule.The most common is zero-order hold (ZOH) defined by A = exp(∆A) and B = (∆A) -1 (exp(∆A)-I)•∆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 a 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,B,C) → (A,B,C), the model can be computed in two ways, either as a 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 a 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,B,C), and consequently (A,B) as 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 as 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 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, 2022; Gu et al., 2022b; Smith et al., 2023) , which we also use.A ∈ R N ×N ,B ∈ R N ×1 ,C ∈ R 1×Nmatrices 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.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. Other closely related SSMs and architectures are discussed further in an extended related work (Appendix B).• Linear attention (Katharopoulos et al., 2020) is an approximation of self-attention (Bahdanau et al., 2014) involving a 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 a standard local convolution, which they frame as a shift-SSM, before the main SSM layer. • Hyena (Poli et al., 2023) uses the same architecture but replaces the S4 layer with an MLPparameterized global convolution (Romero et al., 2021 ). • RetNet (Sun et al., 2023) adds an additional gate to the architecture and simplifies the SSM, allowing an alternative parallelizable computation path, using a variant of attention instead of convolutions. 
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 a technical challenge of how to compute them efficiently.We overcome this with a 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).We argue that a fundamental problem of sequence modeling is compressing context into a smaller state.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 To understand this principle, we focus on 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-aware 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 of LLMs (Olsson et al., 2022) . It requires context-aware 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 an in 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 kernels.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 fundamental 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, a selection mechanism controls how information propagates or interacts along the sequence dimension (see Section 3.5 for more discussion).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 2 illustrates the main selection mechanism that we use. The main difference is simply making several parameters ∆,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 L, meaning that the model has changed from time-invariant to time-varying. This loses the equivalence to convolutions (3) with implications for its efficiency, discussed next.We specifically choose s B (x) = Linear N (x), s C (x) = Linear N (x), s ∆ (x) = Linear D (Linear 1 (x)), and τ ∆ = softplus, where Linear d is a parameterized projection to dimension d. The choice of s ∆ and τ ∆ is due to a connection to RNN gating mechanisms explained in Section 3.5.Hardware-friendly architectures such as convolutions (Krizhevsky et al., 2012) and Transformers (Vaswani et al., 2017) enjoy widespread application. Here we aim to make selective SSMs efficient on modern hardware (GPU) as well. The selection mechanism is quite natural, and earlier works ▷ Time-varying: recurrence (scan) only 7: return y attempted to incorporate special cases of selection, such as letting ∆ vary over time in recurrent SSMs (Gu et al., 2020a) . However, this was computationally difficult, which was why S4 and all 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(BLDN ) FLOPs while the convolutional computation uses O(BLDlog(L)) FLOPs, and the former has a 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 h.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 to a significant speedup compared to a standard implementation.Concretely, instead of preparing the scan input (A,B) of size (B,L,D,N) in GPU HBM (high-bandwidth memory), we load the SSM parameters (∆,A,B,C) directly from slow HBM to fast SRAM, perform the discretization and recurrence in SRAM, and then write the final outputs of size (B,L,D) back to 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) .Finally, we must also avoid saving the intermediate states, which are necessary for backpropagation. We carefully apply the classic technique of recomputation to reduce the memory requirements: the intermediate states are not stored but recomputed in the backward pass when the inputs are loaded from HBM to SRAM. Details of the fused kernel and recomputation are in Appendix D.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 ).This architecture involves expanding the model dimension D by a controllable expansion factor E.For each block, most of the parameters (3ED 2 ) 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 12D 2 parameters of a Transformer's interleaved MHA (multi-head attention) and MLP blocks.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 RNNLinear projectionSequence transformationNonlinearity (activation or multiplication) ⨂ X X X ! X Conv SSM X ! ! Conv SSM 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 modern 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 gate with an activation function. Compared to the MLP block, Mamba adds an SSM to the main branch. For σ we use the SiLU / Swish activation (Hendrycks & Gimpel, 2016; Ramachandran et al., 2017) .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, ∆ in SSMs can be seen to play a generalized role of the RNN gating mechanism. Theorem 1. When N = 1,A = -1,B = 1,s ∆ = Linear(x), and τ ∆ = softplus, then the selective SSM recurrence (Algorithm 2) takes the form g k = σ(Linear(x k )) (the gate) andh k = (1-g k )h k-1 +g k x k .As mentioned in Section 3.2, our specific choices of s ∆ ,τ ∆ is from this connection. In particular, note that if a given input x k should be completely ignored (as necessary in the synthetic tasks), all D channels should ignore it, and so we project the input down to 1 dimension before broadcasting with ∆. Finally, we observe that the function s ∆ (x) = Linear D (Linear 1 (x)) is simply a low-rank projection, and relax the 1 dimension to a larger dimension R, generally set to be a small fraction of D.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 as "um". This property arises because the model can mechanistically filter out any particular input x k , for example in the gated RNN case (Theorem 1) when g k → 0.Filtering context. It 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 E.2.2).We remark that while the A parameter could also be selective, it ultimately affects the model only through its interaction with ∆ via A = exp(∆A) (the discretization). Thus selectivity in ∆ is enough to ensure selectivity in (A,B), and is the main source of improvement. We hypothesize that making A selective in addition to (or instead of) ∆ would have similar performance, and leave it out for simplicity. 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 state-of-the-art 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 (Section 4.5).Table 1 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 LTI SSMs cannot even when combined with more powerful architectures. On Induction Heads, Mamba learns the task perfectly Figure 4 : (Chinchilla scaling laws.) Models of size ≈ 125M to ≈ 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.and can even extrapolate to million-length sequences, or 4000× longer than it saw during training, while no other method goes beyond 2×. Full discussion for synthetic tasks are in Appendix E.1.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, we compare against the standard Transformer architecture (GPT3 architecture), as well as the strongest Transformer recipe we know of (here referred to as Transformer++), 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.Table 2 shows the performance of Mamba on a range of popular downstream 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.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 vocab. It is also known for requiring long-range dependencies to model (Avsec et al., 2021) . We investigate Mamba as a FM backbone for pretraining and fine-tuning in the same setting as recent works on 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 a difficult downstream synthetic classification task requiring long context (Figure 6 ). Full discussion for these tasks is in Appendix E.2. For the audio waveform modality, we compare primarily to the SaShiMi architecture and training protocols (Goel et al., 2022) . The architecture is a UNet with alternating S4 and MLP blocks, which we consider replacing with Mamba. Experiment details are in Appendix F.4.Long-Context Autoregressive Pretraining. We evaluate pretraining quality on YouTubeMix (Deep-Sound, 2017), a standard piano music dataset. Figure 7 evaluates the effect of increasing training sequence lengths from 2 13 = 8192 to 2 20 ≈ 10 6 , while keeping computation fixed. Autoregressive Speech Generation. Tables 3 and 4 show results on the SC09 dataset for generating speech clips of the digits 0-9 (Warden, 2018; Donahue et al., 2019) . Mamba sets significant state-ofthe-art results, and is consistently better than attention or S4 blocks in a parameter-controlled setting.We benchmark the speed of the SSM scan operation (N = 16), as well as the end-to-end inference throughput of Mamba, in Figure 8 . Our efficient SSM scan is faster than the best attention implementation that we know of (FlashAttention-2 (Dao, 2023) ) beyond sequence length 2K, and up to 20-40× faster than a standard scan implementation in PyTorch. Mamba achieves 4-5× higher inference throughput than a Transformer of similar size, since without the KV cache it can use much higher batch sizes. For example, Mamba 6.9B (untrained) has higher inference throughput than a 5× smaller Transformer-1.3B. Details in Appendix F.5.We perform a 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 ). These include (i) combinations of the outer architecture block vs. inner SSM layer (ii) parameterizations of the selective SSM. Results are in Appendix E.3. 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. 
conclusion: 
 We introduce a 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-art results on a diverse set of domains: most notably language, where it matches or 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 as genomics, audio, and video. Our results suggest that Mamba is a strong candidate to be a general sequence model backbone. 
 
 Appendix: Our selection mechanism is inspired by and related to concepts such as gating, hypernetworks, and data-dependence. It can also be viewed as related to "fast weights" (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 (Hochreiter & Schmidhuber, 1997) and GRU (Chung et al., 2014) , or the gated equation in Theorem 1. This was interpreted as a particular mechanism for controlling whether to 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.Hypernetworks 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 a smaller RNN.Data-dependence Similar to hypernetworks, data-dependence can refer to any notion where some parameters of the model depend on the data (Poli et al., 2023) .To illustrate the issues with these concepts, consider a simple diagonal linear layer y = Dx, where D is a diagonal matrix. Now suppose that D is itself generated from a linear transformation of x, with an optional nonlinearity: D = σ(W x). Since it is diagonal, the multiplication becomes an elementwise product:y = σ(W x)•x.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.Thus, while selection mechanisms could be considered a special case of ideas such as architectural gating, hypernetworks, or data-dependence, so can an enormous range of other constructions-essentially anything with a multiplication, including standard attention mechanisms (Bahdanau et al., 2014; 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 (LSTM and GRU), which is a special case (Theorem 1) and also has a deeper history of connections to SSMs through variable (input-dependent) discretization of ∆ (Funahashi & Nakamura, 1993; Tallec & Ollivier, 2018; Gu et al., 2020a) . We also eschew the term "gating" in favor of selection to 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; Kosma et al., 2023) and even attention.We highlight several prior works related to our methods.We describe a brief overview of some structured SSMs from past work, particularly those that have a relation to our method.• 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, 2022) first discovered the empirical effectiveness of diagonal structured SSMs by approximating the HIPPO initialization. This was expanded on theoretically in S4D (Gu et al., 2022b ). • S5 (Smith et al., 2023) independently discovered the diagonal SSM approximation, and is the first S4 model to be computed recurrently with the parallel scan. However, this required lowering the effective state dimension, which they accomplished by switching the SSM dimensions from a SISO (single-input single-output) to MIMO (multi-input multi-output) formulation. • Mega (Ma et al., 2023) introduced a simplification of S4 to be real-instead of complex-valued, giving 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 empirically 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) all focus on the convolutional representation of 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).We use SSM architectures or state space neural networks (SSNN) to refer to deep neural network 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 (Katharopoulos 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 a black box to generate a 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 a selection mechanism (Appendix A.1). 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 et 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 N = 1. This simplification leads to an alternate way to parallelize the computation with a variant of standard multi-head attention instead of convolutions. Although not framed as such, its recurrence can be viewed as a special case of a linear SSM. • 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 MHA and MLP blocks together and was an inspiration for our architecture (Section 3.4) combining the H3 and MLP blocks.RNNs and SSMs are broadly related, as they both involve the concepts of recurrence on a latent state.Several older RNNs such as 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. However, they do not use state expansion (N = 1) or selective B,C parameters, both of which are important for performance (Section 4.6).Additionally, older RNNs famously suffered from efficiency issues and the vanishing gradients problem (Pascanu et al., 2013) , both caused by their sequential nature. The latter could be solved for some of the above RNNs by leveraging the parallel scan (Martin & Cundy, 2018) , but the former 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) ).We also note that there is a 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 & Martínez-Rubio, 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/unitary RNNs are also LTI. 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 (Jing et al., 2019) .Long context has become a 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 a Transformer backbone. It showed ability to generalize up to 1M sequences but only on synthetic memorization tasks; this result is similar to our Induction Heads extrapolation experiment (Figure 3 ). • LongNet (Ding et al., 2023) , which claimed to 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, but did not control for computation time. In fact, its claims about efficiency and performance would be largely matched by any of the LTI S4 variants above.In contrast, we believe this work presents one of the first approaches to meaningfully demonstrate increasing performance with longer context.Proof of Theorem 1. Consider a selective SSM (Algorithm 2) withN = 1, A = -1, B = 1, s ∆ = Linear(x),τ ∆ = softplus. The corresponding continuous-time SSM (1) is h(t) = -h(t)+x(t) which is also called a leaky integrator.The discretization step size is∆ k = τ ∆ (Parameter+s ∆ (x k )) = softplus(Parameter+Linear(x k )) = softplus(Linear(x k ))where we observe that the parameter can be viewed as a learnable bias and folded into the linear projection. Now applying the zero-order hold (ZOH) discretization formulas:A k = exp(∆A) = 1 1+exp(Linear(x k ) = σ(-Linear(x k )) = 1-σ(Linear(x k )) B k = (∆A) -1 (exp(∆A)-I)•∆B = -(exp(∆A)-I) = 1-A = σ(Linear(x k )).Thus the final discrete recurrence (2a) isg k = σ(Linear(x k )) h k = (1-g k )h k-1 +g k x k as desired.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 fusion and recomputation to make SSM scan fast and memory-efficient. We evaluate the speed of our scan implementation compared to convolution and attention in Section 4.5, showing that it is up to 7× times faster than attention at sequence length 32K, and is as memory-efficient as the best attention implementation (FlashAttention).Speed. On modern hardware accelerators (GPUs) most operations (except matrix multiply) are bounded by memory-bandwidth (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,L,D,N ) in GPU HBM (high-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,L,D). However, this requires the number of memory reads/writes on the order of O(BLDN ). We can instead fuse the discretization step, the scan, and the multiplication with C into one kernel:1. We read in O(BLD+DN ) bytes of memory (∆,A,B,C) from slow HBM to fast SRAM.2. We discretize to produce A,B of size (B,L,D,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,L,D) 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 (Section 4.5).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,L,D,N ) 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,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(BLN +DN ), recomputation avoids the cost of reading O(BLN D) elements from HBM. This means that recomputation of the SSM states in the backward pass speeds up the computation compared to storing them and reading them from HBM.Beyond optimizing for the memory requirement of just the scan operation, we also use recomputation 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 a lot of memory but are fast to recompute (e.g. output of activation function or short convolution). As a 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, an each MLP layer stores around 20 bytes of activations per token, for a total of 32 bytes ((assuming mixed-precision 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.1 SYNTHETIC TASKS Full experiment details for these tasks including task details and training protocol are in Appendix F.1.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, LTI 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 as the Denoising task (Jing et al., 2019) .Note that many previous works argue that adding architecture gating (multiplicative interactions) can endow models with "data-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.1).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.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. It requires models to perform associative recall and copy: for example, if the model has seen a 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 a 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 on a range of sequence lengths from 2 6 = 64 up to 2 20 = 1048576 at test time.Models Following established work on induction heads, we use 2 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 a model dimension D of 64 for Mamba and 128 for the other models.Figure 3 shows that Mamba-or more precisely, its selective SSM layer-has the ability to solve the task perfectly because of its ability to selectively remember the relevant token while ignoring everything else in between. It generalizes perfectly to million-length sequences, or 4000× longer than it saw during training, while no other method goes beyond 2×.Out of positional encoding variants for attention models, xPos (which was designed for length extrapolation) is slightly better than the others; also note that all attention models were only tested up to sequence length 2 14 = 16384 due to memory limitations. Out of other SSMs, H3 and Hyena are similar, contrary to the findings in Poli et al. (2023) .For pretraining, we largely follow a 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 a single human genome with about 4.5 billion tokens (DNA base pairs) in the training split. A n Initialization Field Perplexity 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.A n = -1 2 +ni Complex 9.16 A n = -1 2 Real 8.85 A n = -(n+1) Real 8.71 A n ∼ exp(N(0,1)) Real 8.71Selective 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 2 layer models with a model dimension of D = 64.Models are trained for 400K steps at a constant learning rate of 0.0001 with a batch size of 64. We use the Adam optimizer with no weight decay. All models are trained at constant learning rates 2e -4 and 1e -3, and the better results are 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 LR of 2e-4. Mamba learned at both LRs, but extrapolated better at the larger LR of 1e-3.All models were trained on the Pile.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• 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., 2022) and LLaMa (Touvron et al., 2023) . These include:• linear learning rate warmup with cosine decay to 1e-5, with a peak value of 5× the GPT3 value• no linear bias terms• RMSNorm instead of LayerNorm• AdamW hyperparameter β = (.9,.95) (the GPT3 value) instead of the PyTorch default of β = (.9,.999)Architecture and training details. Our models are:• Transformer: The standard Transformer based on GPT3 (Table 11 ).• Transformer++: A Transformer with an improved architecture, namely rotary positional encodings (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 a global convolution parameterized 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.5× 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) a 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 as increasing the learning rates by 2× or 3× 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.This pretraining procedure is the same as the scaling law protocol, but extended to 300B tokens. For the 1.3B model, we use a 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 to 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 EleutherAI (Gao et al., 2021) , as done by most work in this area. We evaluate on the following tasks/datasets that measure common sense reasoning:• LAMBADA (Paperno et al., 2016) .• HellaSwag (Paperno et al., 2016) .• PIQA (Bisk et al., 2020) .• ARC-easy: an easy subset of ARC-challenge.• ARC-challenge (Clark et al., 2018) .• 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). The dataset follows the splits from the prior Enformer work on genomics (Avsec et al., 2021) ; the training split contains a total of S = 34021 segments of length 2 17 = 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 if necessary (e.g. to get longer segments).We deviate from HyenaDNA when the training sequence length is not 2 17 . HyenaDNA always takes a fixed sub-segment (e.g. the beginning or middle of the prescribed segment), and thus for any training The dataset consists of clips of up to 1 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 a multiple of 8 for hardware efficiency, the longest possible sequence is 468×2048 = 958464. The rest of our sequence lengths are defined by successively halving this and rounding up to the nearest multiple of 2048. Training. Models were trained for 200K training steps with a maximum learning rate of 0.002, 20K (10%) warmup steps, and weight decay 0.1 (similar to our general pretraining recipe across domains).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 β = (0.9,0.95)• Gradient clip value 0.1We used a learning rate of 0.002 and 200000 training steps at a batch size of 16.The large Mamba model in Table 3 has 15 layers per stage with an outer dimension of D = 96 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 or NLL. However, automated metrics of generated samples continually improving throughout training.The models in the architecture ablations in Table 4 all have 8 layers per stage with an outer dimension of D = 64 and pooling factor 4. The S4+MLP block has roughly 2D 2 +4D 2 parameters (expansion factor 2 in the MLP). The Transformer block has 4D 2 +2D 2 parameters (expansion factor 1 in the MLP).The Mamba block has the usual ≈ 6D 2 parameters. All models have roughly 6M total parameters.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 PCIe GPU.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 O(Llog(L)) for sequence length L.For attention, we compare against the fastest implementation that we are aware of (FlashAttention-2 (Dao, 2023)), with causal mask. Note that FlashAttention-2 with causal mask is about 1.7× 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 512, 1K, 2K, ... to 512K (some of the baselines run out of memory before reaching 512K). We use a model dimension of D = 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 a 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, 8, 16, 32, 64, to 128, and measure time time taken to generate 128 tokens. We then calculate the throughput (tokens/s) as batch size × 128/time taken. We repeat the measurements 3 times and take the average. Measurements are done on an A100 80GB PCIe GPU. 


Decision: 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.

Reviews: review: paper_summary: paper_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".

While 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. 

Overall, 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.
strength_weakness: strength_weakness: * Lack of results on long-context text tasks to justify that the models can actually replace transformers. 
* 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.
* 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: questions: How good can the model be parallelized and utilize machine compute?

Can you explain why only B, C needed to be context-aware and the static A can still be used?

The 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.
paper_summary: paper_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
performance improves on real data up to million-length sequences. Mamba achieves state-of-the-art performance across
several 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.
strength_weakness: strength_weakness: Don't find a good reason to reject.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.
4. 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.
5. 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.
6. 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}$.

---
[1] Efficient Content-Based Sparse Attention with Routing Transformers (TACL 2020)

[2] Reformer: The Efficient Transformer (ICLR 2020)

[3] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (JMLR 2022)

[4] Sparse Modular Activation for Efficient Sequence Modeling (NeurIPS 2023)

[5] https://github.com/jzhang38/LongMamba

[6] Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024)
questions: questions: 1. How is the training wall time of Mamba compared to other models of similar sizes?
2. How is the performance of Mamba compared to the Llama architecture on downstream NLP tasks?
3. Is the short convolution module critical to Mamba's performance? What if we remove it?
4. 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?
5. Why is Mamba not length extrapolable in real-world language modeling?
6. 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?
paper_summary: paper_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. 

In 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.
strength_weakness: strength_weakness: There is no reason to reject this paper.
questions: questions: * 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.

*  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.

Comments: comment: Thank you for the replies and for sharing interesting future directions! All of them make sense and I would recommend the authors to be clear about those limitations (in language tasks) in the paper itself. I think that would actually make this paper stronger -- instead of leaving an impression of overclaim, it's better to have a deep discussion about the difference, pros and cons of different models. 

As the authors suggested, this direction is still a work-in-progress and a couple of improvements are needed. But still, I think it's a worthy direction and it would be an interesting paper to be published in COLM.