PaperlayerView exact revision on arXiv

Source-linked AI summary

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

arXiv:2312.00752v2cs.LGcs.AI

TL;DR

Transformers dominate sequence modeling, but efficient alternatives have not matched their performance on important modalities such as language. Mamba introduces selective state spaces in an attention-free architecture, achieving Transformer-quality results while scaling linearly with sequence length.

  • Problem

    Existing sequence models face a gap between Transformer-level modeling quality and efficient linear scaling, especially on language and other discrete modalities.

  • Method

    Mamba makes SSM parameters input-dependent so it can selectively retain or forget information, combining selective state spaces in a simplified attention-free architecture.

  • Results

    Mamba achieves state-of-the-art performance across language, audio, and genomics, matching or exceeding strong Transformer models while improving with context lengths up to 1M.

  • Takeaways & Limitations

    Selective state spaces make Mamba a strong candidate for general sequence-model backbones across domains requiring long context.

  • Takeaways & Limitations

    The empirical evaluation uses small models, so it remains unclear whether Mamba compares favorably at larger parameter scales.

Abstract

from arXiv · show

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 (5$\times$ 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.

1 Introduction

Mamba introduces selective state spaces whose input-dependent parameters enable content-based filtering while retaining linear sequence scaling. Across synthetic, audio, genomics, and language tasks, it achieves strong quality with efficient long-context computation.

  • Selection Mechanism: Input-dependent SSM parameters let Mamba selectively filter irrelevant information and retain relevant information indefinitely, addressing prior models’ weakness on discrete, information-dense data.The mechanism is motivated by selective-copy and induction-head tasks and enables content-based selection.
  • Hardware-aware Algorithm: A hardware-aware recurrent scan computes selective SSMs without materializing the expanded state, avoiding costly GPU-memory-hierarchy transfers.This change replaces the convolutional computation used by time- and input-invariant prior SSMs.
  • Architecture: Mamba combines prior SSM designs with a Transformer MLP block into a single homogeneous architecture incorporating selective state spaces.The resulting model is fully recurrent, with computation and memory scaling linearly in sequence length during training.
  • Empirical validation: On synthetic copying and induction-head tasks, Mamba solves the problems easily and extrapolates solutions beyond 1M tokens.These tasks are presented as important tests proposed to capture capabilities relevant to large language models.
  • Empirical validation: Mamba outperforms prior state-of-the-art models on audio waveforms and DNA sequences, while improving with longer contexts up to million-length sequences.The comparisons include SaShiMi, Hyena, and Transformers, with downstream and pretraining gains; speech-generation FID is reduced by more than half.
  • Empirical validation: Mamba is the first linear-time sequence model reported to achieve Transformer-quality language performance in pretraining perplexity and downstream evaluations, with 5× generation throughput.Scaling experiments up to 1B parameters exceed a broad range of baselines, including modern LLaMA-based Transformer training recipes.

2 State Space Models

Structured SSMs are latent-state sequence transformations that can be computed as either recurrences or convolutions, with efficiency enabled by time-invariant dynamics and structured state matrices. This section identifies linear time invariance as a limitation for content-based modeling and motivates removing it while preserving efficient computation.

  • Structured SSMs: SSMs map sequences through higher-dimensional latent states and use four parameters, (Δ, A, B, C), to define the transformation.Structured SSMs apply the model independently across input channels.
  • Computation: After discretization, SSMs support both linear-recurrence and global-convolution computation, typically using convolution for parallel training and recurrence for autoregressive inference.The two modes exploit different access patterns: full sequences during training versus one timestep at a time during inference.
  • Linear Time Invariance: LTI SSMs keep (Δ, A, B, C) fixed across time, making recurrence and convolution equivalent but limiting their ability to model certain data types.The work therefore targets removing the LTI constraint without reintroducing the associated efficiency bottlenecks.
  • Structure and Dimensions: Diagonal structure makes the state matrices compact, but applying an SSM over batch size B, length L, and D channels still requires O(BLDN) time and memory.The total hidden state has dimension DN per input, creating the fundamental efficiency bottleneck addressed later.
  • SSM Architectures: SSMs are standalone sequence transformations that can be embedded in end-to-end architectures, including related designs such as linear attention, H3, Hyena, RetNet, and RWKV.These architectures use recurrent, convolutional, gated, or attention-like variants of SSM computation.

3 Selective State Space Models

Mamba introduces selective state space models whose input-dependent parameters enable content-aware filtering and propagation of sequence information. The approach replaces convolution with hardware-aware recurrent computation and combines selective SSMs into a simplified architecture without attention or MLP blocks.

  • 3.1 Motivation: LTI models fail on selective copying and induction-head tasks because constant dynamics cannot select or update hidden states according to input content.Selective copying requires memorizing relevant tokens amid irrelevant ones, while induction heads require context-aware retrieval.
  • 3.1 Motivation: Selective SSMs make sequence interactions input-dependent, allowing the model to focus on relevant inputs and filter them into a sequential state.This addresses the efficiency–effectiveness tradeoff: efficient models compress context, while effective models retain necessary information.
  • 3.2 Selective State Space Models: The S6 mechanism makes Δ, B, and C functions of the input, changing the SSM from time-invariant to time-varying and restricting computation to recurrent scans.This removes the equivalence to convolutions while enabling input-dependent state transitions.
  • 3.3 Hardware-Aware Algorithms: Kernel fusion, parallel scan, and recomputation make selective recurrent SSMs efficient on modern hardware despite losing convolutional computation.The naive recurrent form uses O(BLDN) FLOPs, while convolution uses O(BLDlog(L)) FLOPs; recurrent computation can use fewer FLOPs for long sequences and moderate N.
  • 3.4 Mamba: Mamba simplifies prior SSM architectures by combining the linear-attention-inspired and MLP components into one homogeneous stacked block without attention or separate MLP blocks.Selective models can also reset their state to remove irrelevant history from longer contexts.

4 Empirical Evaluation

Across synthetic tasks and language, DNA, and audio evaluations, Mamba’s selective SSM enables content-based sequence reasoning, strong scaling, long-context use, and competitive generation quality. Ablations further show that selective SSMs are the main source of performance gains.

  • 4.1 Synthetic Tasks: Mamba’s selective SSM perfectly solves induction-heads and generalizes to million-length sequences, 4000× longer than training, while no other method exceeds 2×.It selectively remembers relevant tokens while ignoring intervening information.
  • 4.2 Language Modeling: Mamba is the first attention-free model to match the strong Transformer++ recipe, with scaling advantages especially as sequence length grows.The comparison covers models from approximately 125M to 1.3B parameters under the Chinchilla protocol.
  • 4.3 DNA Modeling: At approximately 40M parameters, Mamba matches Transformer++ and HyenaDNA with roughly 3× to 4× fewer parameters and improves smoothly with model size.Mamba also benefits from context lengths up to 1M, whereas HyenaDNA degrades as sequence length increases.
  • 4.4 Audio Modeling: A small Mamba-UNet outperforms larger GAN- and diffusion-based speech models, while a parameter-matched larger model further improves fidelity metrics dramatically.On SC09, Mamba also consistently outperforms S4+MLP in outer blocks and ranks above S4+MLP and MHA+MLP in center blocks.
  • 4.6 Ablations: Selective SSMs substantially improve performance over static alternatives, while Mamba performs similarly to H3 and is slightly better when using a selective layer.Changing the Δ and (B, C) projections from static to selective provides the largest benefit.
  • 4.6 Ablations: Increasing the selective SSM state size yields over a 1.0 perplexity improvement for only 1% additional parameters.Larger projection dimensions generally provide more modest gains with small parameter increases.

5 Discussion

The discussion shows that selective SSM design choices improve expressivity and state capacity, while highlighting tradeoffs across modalities, unresolved downstream capabilities, and uncertainty at larger scales.

  • Ablations: Expressivity of Δ: Selective Δ projections improve performance even at dimension 1, with further gains from larger projections but a modest parameter cost.The SSM state size is fixed at N = 16.
  • Ablations: SSM state dimension: Increasing the SSM state dimension N significantly improves performance at negligible parameter and FLOP cost only when B and C are selective.The comparison includes both constant and selective B and C, with the Δ projection fixed to 64.
  • No Free Lunch: Continuous-Discrete Spectrum: The selection mechanism helps SSMs overcome weaknesses on discrete modalities such as text and DNA but can impede performance on data where LTI SSMs excel.This reflects a continuous–discrete spectrum tradeoff discussed alongside audio waveform ablations.
  • Downstream Affordances: Whether SSMs match Transformer-based foundation models in downstream affordances such as fine-tuning, prompting, in-context learning, and RLHF remains an open question.The discussion also mentions adaptation, instruction tuning, and quantization as relevant pretrained-model interactions.
  • Scaling: Mamba’s evaluation used only small models, leaving its competitiveness, scaling behavior, and engineering requirements at larger sizes unresolved.The discussion contrasts this limitation with prior evaluations of Llama, RWKV, and RetNet at 7B parameters and beyond.

6 Conclusion

The paper introduces selective state space models that enable context-dependent reasoning while retaining linear sequence-length scaling. In an attention-free architecture, Mamba achieves state-of-the-art performance across diverse domains, matching or exceeding strong Transformers and motivating broader foundation-model applications.

  • Conclusion: Selective state space models enable context-dependent reasoning while scaling linearly with sequence length.The selection mechanism is the paper’s central contribution to structured state space models.
  • Conclusion: Mamba’s simple attention-free architecture achieves state-of-the-art results across diverse domains, matching or exceeding strong Transformer models.This establishes the approach as a competitive general-purpose sequence-modeling architecture.
  • Conclusion: Selective state space models could support foundation models across different domains.The conclusion highlights their potential for broad applications beyond the demonstrated settings.

A Discussion: Selection Mechanism · B Related Work · B.1 S4 Variants and Derivatives

The paper distinguishes selection from broad labels such as gating, hypernetworks, and data-dependence, grounding it instead in RNN gating and input-dependent SSM discretization. It situates S6 among S4-derived models, emphasizing its selective, SISO recurrent formulation against prior largely non-selective, often LTI structured SSMs.

  • A Discussion: Selection Mechanism: Selection is distinct from broad gating, hypernetwork, and data-dependent labels because those categories include trivial multiplicative constructions and standard attention mechanisms.A GLU-like transformation can satisfy all three labels while functioning only as an activation, making those labels uninformative for characterizing selection.
  • A Discussion: Selection Mechanism: Selection is most closely related to traditional RNN gating and connects to SSMs through input-dependent discretization of Δ (Funahashi and Nakamura 1993; Gu, Dao, et al. 2020; Tallec and Ollivier 2018).The paper uses “selection” rather than “gating” to avoid the latter’s overloaded modern usage.
  • B Related Work: The related-work overview places Mamba alongside recurrent layers S4, S5, and quasi-RNNs, and end-to-end architectures H3, RetNet, and RWKV.These models are identified as among the methods most closely related to the paper’s approach.
  • B.1 S4 Variants and Derivatives: S4 introduced structured SSMs with diagonal and diagonal-plus-low-rank structure, using efficient convolutional algorithms motivated by HIPPO’s connection to continuous-time online memorization (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021).DSS later established the empirical effectiveness of diagonal structured SSMs, while S4D expanded this result theoretically (Gupta, Gu, and Berant 2022; Gu, Gupta, et al. 2022).
  • B.1 S4 Variants and Derivatives: S5 was the first S4 model computed recurrently with a parallel scan, but reduced effective state dimension by switching from SISO to MIMO; S6 retains SISO dimensions while sharing the scan (Smith, Warrington, and Linderman 2023).The supplied passage identifies S6’s SISO formulation and shared scan as its distinctions from S5.
  • B.1 S4 Variants and Derivatives: Lu et al. (2023) use a hard-coded reset mechanism for S5 in meta-RL, whereas the paper characterizes selection as a learnable, input-dependent alternative.Their mechanism manually sets A to 0 at episode boundaries; the paper proposes probing whether selective SSMs can learn such resets automatically.
  • B.1 S4 Variants and Derivatives: Most other structured SSM variants remain non-selective and usually strictly LTI, including convolutional methods that cannot perform fast autoregressive inference directly.Liquid S4 adds a limited input-dependent transition while remaining close to LTI, whereas SGConv, Hyena, LongConv, MultiresConv, and Toeplitz Neural Network focus on long convolution kernels.

B.2 SSM Architectures … C Mechanics of Selective SSMs

The paper situates selective SSMs among gated SSM, recurrent, linear-attention, and long-context architectures, emphasizing input-dependent selection and principled SSM parameterization. It also formalizes selective SSM mechanics through discretization, while contrasting its long-context validation with prior claims.

  • B.2 SSM Architectures: Selective SSM architectures are positioned against GSS, Mega, H3, Selective S4, RetNet, RWKV, and GAU, differing in how they combine recurrence, attention-like mechanisms, gating, and state expansion.GSS contracts model dimensions, whereas the paper’s architecture expands them to increase SSM state size; GAU also inspires the architecture.
  • B.2 SSM Architectures: Selective S4 (J. Wang et al. 2023) masks inputs externally, which the paper distinguishes from selection because masking does not alter the spacing between relevant tokens in Selective Copying.The paper therefore characterizes Selective S4 as architectural gating rather than a selection mechanism.
  • B.3 Relationship to RNNs: RNNs and SSMs both use latent-state recurrence, while older RNNs faced sequential-efficiency and vanishing-gradient problems that later SSM theory addressed through parallelization and careful parameterization.Orthogonal and unitary RNNs instead constrain transition matrices to control eigenvalues and prevent vanishing gradients, but retain other limitations.
  • B.3 Relationship to RNNs: Gated RNNs such as strongly typed RNNs, QRNNs, and SRUs can be viewed as selective SSMs, but lack state expansion and selective B,C parameters important for performance.The paper further contrasts their heuristic gates with a selection-plus-discretization formulation grounded in SSM theory.
  • B.4 Linear Attention: Linear Attention (Katharopoulos et al. 2020) connects kernel attention with recurrent autoregressive models, motivating numerous kernel, normalization, positional, and sampling-based variants.The section also points to Tay, Dehghani, Bahri, et al. (2022) for a broad survey of efficient-attention variants.
  • B.5 Long Context Models: The paper argues that it is among the first approaches to meaningfully demonstrate improving performance with longer context, unlike prior long-context claims limited by synthetic tasks, shorter evaluations, or confounded data and compute.Examples include Recurrent Memory Transformer, LongNet, Hyena and HyenaDNA, and Sparse Transformer, whose reported context capabilities were not comparably validated.
  • C Mechanics of Selective SSMs: The selective SSM mechanics proof specializes Algorithm 2 to N=1, A=-1, B=1, with input-dependent sΔ=Linear(x) and τΔ=softplus, then applies zero-order-hold discretization to obtain the final discrete recurrence.The discretization also allows a parameter to be interpreted as a learnable bias folded into the linear projection.

D Hardware-aware Algorithm For Selective SSMs … E.2 Language Modeling

Mamba’s selective SSMs use a hardware-aware fused scan to preserve linear scaling while substantially reducing memory traffic and runtime, with recomputation keeping activation memory competitive with optimized Transformers. The supplied experimental details describe synthetic-task settings and training protocols for selective copying and induction-head evaluations.

  • D Hardware-aware Algorithm For Selective SSMs: Selective SSM scans remain linear in sequence length, with theoretical cost O(BLDN), but require hardware-aware implementation for foundation-model training.Input-dependent selectivity removes the convolution and FFT implementation available to nonselective SSMs, so the method uses a parallel associative scan.
  • D Hardware-aware Algorithm For Selective SSMs: 20-40× speedups result from fusing discretization, associative scanning, and output multiplication, reducing memory I/O by O(N).The fused kernel keeps intermediate computations in SRAM and writes only final outputs to HBM.
  • D Hardware-aware Algorithm For Selective SSMs: For sequences too long to fit in SRAM, chunked fused scans continue across chunks using intermediate scan states.This extends the fused implementation beyond SRAM capacity while preserving the scan recurrence.
  • D Hardware-aware Algorithm For Selective SSMs: 16 bytes of activation memory per token for a selective SSM makes two selective-SSM layers comparable to one FlashAttention layer plus one MLP layer at 32 bytes per token.The comparison assumes mixed-precision FP16 or BF16 training.
  • E.1 Synthetic Tasks: Synthetic-task experiments use length-4096 sequences, 16-token vocabularies, 16 memorized data tokens, two layers, and model dimension D=64.The vocabulary includes a white noise token from Figure 2.
  • Selective Copying.: Selective-copying evaluations use randomly generated training data, fixed randomly generated validation sets, and model-specific reporting checkpoints ranging from 10 to 50 epochs.Mamba and MHA-Abs results are reported after 25 epochs, MHA-RoPE and MHA-xPos after 50, and LTI H3 and Hyena after 10.
  • Selective Copying.: Selective-copying training uses Adam without weight decay, constant learning rates 2e−4 and 1e−3, and reports the better result for each model.Attention and Hyena did not learn at 1e−3; H3 generalized better at 2e−4, while Mamba extrapolated better at 1e−3.

E.2.1 Scaling Law Details

Scaling-law experiments trained models on the Pile with the GPT2 tokenizer, largely following the GPT3 recipe while adjusting batch size and token budgets. The study compared standard and improved-recipe versions of Transformer, Hyena, H3++, RWKV, RetNet, and Mamba.

  • Scaling-law experiments trained all models on the Pile with the GPT2 tokenizer, generally following the GPT3 recipe.
  • Model Sizes: The scaling-law model sizes and hyperparameters followed GPT3 specifications with minor modifications to the 1.3B batch size and training budgets.The 1.3B batch size changed from 1M to 0.5M tokens, while training steps and total tokens were adjusted toward Chinchilla scaling laws.
  • Training Recipes: The default training recipe used AdamW, gradient clipping at 1.0, linear warmup with cosine decay, and the GPT3 peak learning rate.
  • Training Recipes: An improved recipe raised the peak learning rate to 5× GPT3’s value, decayed to 1e−5, removed linear biases, used RMSNorm, and changed AdamW β to (.9, .95).These changes were inspired by recipes used in PaLM (Chowdhery et al. 2023) and LLaMa (Touvron et al. 2023).
  • Architecture and Training Details: The comparison included GPT3-style Transformer, improved Transformer++, Hyena, H3++, RWKV, RetNet, and standard Mamba architectures.Transformer++ used rotary positional encodings and SwiGLU; Hyena and H3++ used specified architectural modifications, while Mamba used the improved recipe.

E.2.2 Additional Scaling Law Ablations

Additional ablations show that interleaving Mamba with MLP or attention blocks changes performance only slightly, while training recipes and expanded SSM state dimensions have larger effects. The choice of inner LTI SSM has little impact.

  • Mamba Architecture: Interleaving Blocks: Mamba-MLP is only slightly worse than homogeneous Mamba and remains better than all models except Transformer++, while Mamba-MHA is only slightly better.The ablations test interleaving Mamba blocks with standard MLP or multi-head attention blocks instead of stacking Mamba blocks homogeneously.
  • H3 Architecture: Training Recipes: A large improvement comes from the improved training recipe, while replacing the inner LTI SSM has little effect; expanding head dimension improves performance.These findings support the importance of training choices and expanded state dimension for SSM performance, whereas Hyena and S4 choices are comparatively unimportant.

E.2.3 Downstream Evaluation Details … E.4 Audio Details

The appendix details evaluation protocols for language, genomics, and long-context classification, including dataset construction, model scaling, optimization, and sequence-length schedules. It also reports Great Apes DNA classification results using pretrained models matched to each context length.

  • E.2.3 Downstream Evaluation Details: Downstream language evaluation uses EleutherAI’s LM evaluation harness on LAMBADA, HellaSwag, PIQA, ARC-challenge, ARC-easy, and WinoGrande.Accuracy is reported for LAMBADA, WinoGrande, PIQA, and ARC-easy, while HellaSwag and ARC-challenge use accuracy normalized by sequence length.
  • E.3.1 Pretraining Details: The HG38 training set follows Enformer’s genome splits, containing 34021 segments of length 2^17 = 131072 and approximately 4.5 billion DNA tokens.Segments are represented by chromosome number and start and end indices, and can be extended for longer sequences.
  • E.3.1 Pretraining Details: For context length L <= 2^17, each genome segment is divided into non-overlapping length-L subsegments, whereas longer contexts produce two samples per segment anchored at its prescribed boundaries.At length 2^18, training uses 4× as many tokens as the default, and at length 2^20 it uses 16× as many.
  • E.3.2 Scaling: Model Size Details: Scaling compares Transformer++, HyenaDNA, and standard Mamba, doubling Mamba’s block count so parameter counts match Transformer layers containing both attention and MLP blocks.Transformer++ uses RoPE positional encodings, while HyenaDNA replaces multi-head attention with an H3 global-convolution block parameterized by an MLP.
  • E.3.2 Scaling: Model Size Details: Learning-rate sweeps selected 2e-3 for Transformer and HyenaDNA and 8e-3 for Mamba, with Mamba more stable and improving further at the higher rate.The sweep used {1e-3, 2e-3, 4e-3, 8e-3}; Mamba also outperformed baselines at the matched 2e-3 rate.
  • E.3.3 Scaling: Context Length Details: Context-length scaling uses sequence-length warmup for two epochs at each power-of-two length from 2^10 = 1024, while halving batch size whenever sequence length doubles.The schedule was not tuned, and its benefit for DNA pretraining was not tested by disabling it.
  • E.3.4 Species (Great Apes) Classification: Species classification uses only the final causal output for prediction while holding batch_size × sequence_length constant during pretraining, and applies additional warmup for contexts above 2^14 = 16384.A 2^20 = 1048576-context model uses six warmup epochs followed by four epochs at maximum length.
  • E.3.4 Species (Great Apes) Classification: Great Apes DNA classification is evaluated from 2^10 = 1024 through 2^20 = 1048576 context lengths using pretrained models with matching context lengths.Table 13 reports accuracy, with random guessing at 20%.

E.4.1 YouTubeMix Audio Pretraining · E.4.2 SC09 Speech Generation

YouTubeMix experiments show that audio waveforms favor LTI processing near the raw signal, while selection becomes less harmful after tokenization and compression. SC09 training used autoregressive settings, but the small dataset caused overfitting in likelihood metrics for the largest model.

  • E.4.1 YouTubeMix Audio Pretraining: YouTubeMix models used 15 total Mamba blocks, pooling factor p=16, outer dimension D=64, and approximately 3.5M parameters.The dataset used 8-bit mu-law encoding with a vocabulary of 256 discrete tokens.
  • E.4.1 YouTubeMix Audio Pretraining: YouTubeMix clips lasted up to 1 minute or length 960000, with the longest hardware-efficient sequence set to 468 × 2048 = 958464.Longer clips were subsampled and divided into desired sequence lengths after two pooling stages.
  • E.4.1 YouTubeMix Audio Pretraining: Audio waveforms significantly favor Mamba-S4 over selection-based Mamba-S6, consistent with their smooth, uniformly sampled continuous structure.Figure 10 reports that removing selection improves long-form audio waveform pretraining because LTI models provide a matching inductive bias.
  • E.4.1 YouTubeMix Audio Pretraining: Ablating selection only from inner U-Net layers sharply reduces performance differences, suggesting that layers after tokenization need not remain LTI.The outer layers remain Mamba-S4, while inner layers operate after the signal has been tokenized and compressed.
  • E.4.1 YouTubeMix Audio Pretraining: The real-valued SSM still underperforms the complex-valued SSM when selection is removed from the inner layers.This comparison is reported in the setting where outer U-Net layers remain Mamba-S4.
  • E.4.2 SC09 Speech Generation: SC09 autoregressive training used 0.002 learning rate, 200000 steps, batch size 16, 10% warmup, and gradient clipping at 0.1.These settings largely followed the autoregressive language-modeling protocol.
  • E.4.2 SC09 Speech Generation: The large SC09 Mamba model used 15 layers per stage, D=96, and pooling factor 4, but the small dataset caused substantial BPB or NLL overfitting.Automated metrics for generated samples continued improving throughout training despite the likelihood overfitting.
  • E.4.2 SC09 Speech Generation: SC09 architecture ablations used 8 layers per stage, D=64, pooling factor 4, and roughly 6M parameters across S4+MLP, Transformer, and Mamba blocks.The S4+MLP, Transformer, and Mamba blocks were parameterized to comparable total model sizes.

E.5 Efficiency Benchmark

The efficiency benchmark evaluates selective SSM scans against convolution and attention, alongside end-to-end inference and memory measurements. Mamba’s memory footprint is comparable to the most optimized Transformer, while its scan implementation fuses discretization with scanning to avoid materializing large parameters in HBM.

  • Scan Operation: The core-operation benchmark measures selective-SSM scans against convolution and attention on an A100 80GB PCIe GPU, excluding costs outside each core operation.Convolution uses a PyTorch FFT implementation, while attention uses causal FlashAttention-2 (Dao 2024).
  • Scan Operation: The fused selective-SSM scan combines discretization and parallel scanning, avoiding materialization of the large A, B, and C parameters in HBM.The unfused PyTorch baseline materializes these parameters in HBM, whereas the proposed implementation fuses the operations.
  • Memory Benchmark: Mamba’s memory footprint is comparable to the most optimized Transformer for 125M models trained on sequences of length 2048.The benchmark compares training memory on one A100 80GB GPU against a Transformer using torch.compile kernel fusion and FlashAttention-2.
  • End-to-end Inference: End-to-end inference compares Mamba 1.4B and 6.9B models with 1.3B and 6.7B GPT-3-style Transformers across batch sizes from 1 to 128.Measurements generate 128 tokens from 2048-token prompts on an A100 GPU, averaging three runs.
Loading 2312.00752v2…