Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

State Space Duality

• SSD connects structured SSMs and attention through structured matrices, equates SSMs with semiseparable matrices, and reframes SSM computation methods as structured-matrix multiplication algorithms.The paper presents matrix-transformation and tensor-contraction views of the duality.Source: §1 · p. 2, Figure 1 · p. 2
• The paper proves linear attention’s recurrent form through tensor contractions, generalizes it to structured masked attention, and shows that any kernel attention method with a fast recurrent form is an SSM.SSMs and structured masked attention have a large intersection whose members admit both linear SSM and quadratic attention forms.Source: §1 · p. 2

Efficient Algorithms

• The block-decomposed SSD algorithm combines linear SSM recurrence with a quadratic dual form, and its implementation runs 2–8× faster than optimized Mamba selective scan.It crosses over with optimized FlashAttention-2 at sequence length 2K and is 6× faster at sequence length 16K.Source: §1 · p. 2

Architecture Design

• The architecture framework treats Mamba as a multi-input SSM analogous to multi-value attention and introduces head-structure choices that inform Mamba-2.Mamba-2 combines a modified parallel Mamba block with SSD as its inner SSM layer; the modifications include grouped-value attention heads and parallel data-dependent projections.Source: §1 · p. 2
• Mamba-2 Pareto-dominates Mamba and Transformer++ in Chinchilla scaling laws, and 2.7B Mamba-2 trained on 300B Pile tokens outperforms Mamba-2.8B, Pythia-2.8B, and Pythia-6.9B.The scaling-law comparison uses the same setting as Mamba; the downstream comparison uses models trained on the same dataset.Source: §1 · pp. 2–3

Systems Optimizations

• The framework extends Transformer systems optimizations to SSMs, including tensor and sequence parallelism and padding-free variable-length Mamba-2 training.Sequence parallelism passes recurrent states between devices.Source: §1 · p. 3

2.1 Structured State Space Models

• SSMs use linear recurrences with state size N as an expansion factor; time-invariant dynamics yield global convolutions, while selective Mamba varies parameters to choose inputs.Selective SSMs are computed recurrently rather than convolutionally and do not leverage accelerator matrix-multiplication units.Source: §2.1 · p. 4

2.2 Attention

• Softmax attention compares every pair of sequence positions, causing quadratic training cost; linear attention drops softmax and uses associativity to rewrite the computation, with causal masks yielding recurrences.The paper’s structured masked attention formulation generalizes these linear-attention ideas.Source: §2.2 · p. 5

2.3 Structured Matrices

• Structured matrices support compressed representations and fast operations; the paper connects SSMs to a class not previously used in deep learning and derives efficient methods from that connection.Source: §2.3 · p. 5

2.4 Overview: Structured State Space Duality

• SSD simplifies selective SSM dynamics to scalar-times-identity structure and uses larger head dimensions, trading some expressive power for training efficiency and accelerator matrix-unit use.Typical head dimensions are P = {64, 128}.Source: §2.4 · p. 5
• SSD’s quadratic dual is attention-like: it drops softmax and applies a data-dependent mask L that controls information transfer across time.This is a structured masked attention generalization of linear attention.Source: §2.4 · p. 6
• Block decomposition of the SSD matrix enables hardware-efficient multiplication with better efficiency tradeoffs than the pure linear or quadratic forms.The algorithm uses the shared matrix representation of SSD, in which the recurrent form is linear-time and the dual form is quadratic-time.Source: §2.4 · p. 6

3.1 The Matrix Transformation Form of State Space Models

• The matrix-transformation view writes an SSM as a matrix multiplication mapping an input sequence x to an output sequence y.The section develops this view to study SSMs as sequence transformations and establish their connection to semiseparable matrices.Source: §3.1 · p. 7, §3 · p. 7

3.2.1 The Sequentially Semiseparable (SSS) Representation

• A lower triangular matrix is N-semiseparable when each submatrix on or below its diagonal has rank at most N; this class is exactly equivalent to SSS representations.The SSS representation can be compressed to O(NT) parameters, which is asymptotically tight (Pernet, Signargout, and Villard 2023).Source: §3.2 · p. 7, §3.2.1 · p. 8

3.2.2 1-Semiseparable Matrices: the Scalar SSM Recurrence

• The 1-semiseparable matrix multiplication is the scalar SSM recurrence, or cumprodsum, and can be computed through the recurrence y_t = a_t y_{t-1} + x_t.The paper identifies fast scalar-recurrence algorithms with structured factorizations of 1-semiseparable matrices.Source: §3.2.2 · pp. 8–9

3.3 State Space Models are Semiseparable Matrices

• An SSM transformation with state size N is exactly multiplication by an N-semiseparable matrix in sequentially semiseparable representation.Writing the sequence transformation as a matrix multiplication reveals this equivalence.Source: §3.3 · p. 9

3.4 Computing State Space Models through Structured Matrix Algorithms

• O(NT) parameters represent an N-semiseparable matrix, and its matrix-vector multiplication takes O(NT) time and space.The matrix perspective reduces efficient SSM computation to algorithms for structured matrix multiplication.Source: §3.4 · p. 10

3.4.1 The Linear (Recurrent) Mode

• O(TN) time computes any state space model of state size N on sequence length T, excluding potential preprocessing.The result applies to all real SSMs, including dense unstructured A matrices, and matches the O(TN) lower bound from the sizes of B and C; preprocessing may be needed.Source: §3.4.1 · pp. 10–11

3.4.2 The Quadratic (Naive) Mode

• Materializing the SSM transformation matrix costs quadratic time in sequence length, but can be more efficient for short sequences because of constant factors and hardware-friendly matrix multiplications.For a particular structured-SSM case, this computation resembles quadratic attention.Source: §3.4.2 · p. 11

3.4.3 Summary

• The semiseparable-matrix perspective views forward-pass methods as matrix-multiplication algorithms, with SSD's dual modes corresponding to linear-time structured and quadratic-time naive multiplication.Exploiting semiseparable structure can enable better algorithms and further insights.Source: §3.4.3 · p. 11

4.1.1 Attention

• Attention maps query, key, and value sequences to an output, and the general formulation separates source, target, feature, and head dimensions.Source: §4.1.1 · p. 12

4.1.2 Self-Attention

• The general attention formulation abstracts away self-attention’s equal sequence and feature dimensions, allowing the framework to capture variants such as cross-attention.Source: §4.1.2 · p. 12

4.1.3 Kernel Attention

• Kernel attention views exponentiated query-key products as feature-map inner products, abstracting the map into the queries and keys while setting aside normalization.The normalization can be handled by passing in V = 1 and dividing.Source: §4.1.3 · p. 12

4.1.4 Masked (Kernel) Attention

• Masked kernel attention applies a mask matrix that can enforce causality or encode sparse patterns such as banded, dilated, or block-diagonal structure.Source: §4.1.4 · p. 13

4.2 Linear Attention

• Autoregressive kernel attention with a causal mask can be computed in O(T) time by a recurrence taking constant time per step.Source: §4.2 · p. 14

4.2.1 A Tensor Contraction Proof of Linear Attention

• For the causal mask, the proof identifies multiplication by the mask as a feature-wise cumulative sum, avoiding the sequence-length-quadratic bottleneck.The alternate contraction order gives a rigorous derivation of linear attention.Source: §4.2.1 · p. 14

4.3 Structured Masked Attention

• Structured masked attention extends linear attention to any mask matrix with sub-quadratic matrix-vector multiplication, which can make the linear-mode computation sub-quadratic.Its linear-mode algorithm optimizes the mask multiplication step; semiseparable structured attention is the focus of the following section.Source: §4.3 · pp. 14–15

4.3.1 Summary: The Dual Forms of Masked Attention

• Linear attention’s quadratic and linear algorithms compute the same masked-attention function using different contraction orders, producing different computational complexities.The paper connects this attention duality to the dual quadratic and linear forms of state-space models.Source: §4.3.1 · p. 15

5.1 Scalar-Identity Structured State Space Models

• Materializing a scalar structured SSM’s semiseparable matrix and multiplying it quadratically is exactly masked kernel attention.This equivalence holds for the scalar-structure case of the SSM.Source: §5.1 · p. 16

5.2 1-Semiseparable Structured Masked Attention

• 1-semiseparable structured masked attention generalizes linear attention with a scalar-SSM recurrence and is a diagonal SSM when A is scalar-identity.Its recurrence replaces linear attention’s cumsum with a scalar SSM scan.Source: §5.2 · pp. 16–17
• Any structured masked attention that is autoregressive with bounded order must use a semiseparable mask, characterizing efficient autoregressive attention.The section identifies general semiseparable structured masked attention as the class of efficient autoregressive attention.Source: §5.2 · p. 17
• General semiseparable structured masked attention is more expressive than 1-semiseparable attention and cannot be represented by a standard SSM.The semiseparable multiplication and linear attention form can nevertheless be absorbed into a similar 1-semiseparable instance using a single, larger expansion.Source: §5.2 · p. 17

5.3 Structured State-Space Duality (SSD)

• Scalar-identity SSMs and 1-semiseparable structured masked attention are duals with the same linear and quadratic forms.SSMs are usually expressed as linear-time recurrences but can be expanded into quadratic forms; attention’s quadratic interactions can be reduced to a linear form by changing the contraction order.Source: §5.3 · p. 17

6 A Hardware-Efficient Algorithm for SSD Models

• O(TN^2) training FLOPs, O(TN) inference FLOPs, and O(N^2) inference memory suffice when N=P, with work dominated by matrix multiplications.These bounds are tight: the state has size N^2 and the input has TN elements.Source: §6 · pp. 17–18

6.2 Low-Rank Blocks

• Off-diagonal blocks are low-rank factorizations: right factors form chunk states, a scalar SSM scan propagates states across chunks, and left factors produce outputs from prior inputs.The right-factor step assumes no inputs before the chunk; the left-factor step uses the correct initial state and excludes inputs within the chunk.Source: §6 · p. 18, §6.2 · p. 20

6.1 Diagonal Blocks

• For small chunk lengths Q, diagonal blocks use the dual quadratic SMA form to compute within-chunk outputs, and the chunks can be computed in parallel.Each diagonal-block subproblem accounts only for inputs within its chunk, assuming a zero initial state.Source: §6.1 · p. 19

6.3 Computational Cost

• With N=P=Q, total work is O(TN^2) FLOPs and O(TN) memory, mostly N×N matrix multiplications; the scalar scan is negligible on GPU and Q times cheaper than a pure SSM scan.The scan costs O(TN^2/Q) FLOPs; blocking also allows other scan algorithms that may be easier to implement without significant slowdown.Source: §6.3 · p. 22

7.1 Block Design

• Mamba-2 produces A, X, B, and C in parallel at block start, adds normalization before the final output projection, and shares a single B, C head across X heads.The extra normalization was added to alleviate instabilities observed in larger models.Source: Figure 6 · p. 23, §7.1 · p. 23

7.2 Multihead Patterns for Sequence Transformations

• The framework transfers MHA-style head patterns to SSMs, typically fixing state and head dimensions near 64 or 128 as model width grows by adding heads.MCS shares X and B across heads but gives C independent copies; MES shares C and X while B varies by head.Source: §7.2 · p. 24
• Among controlled head-pattern variants with matched parameter counts and total state dimension, MVA—the original Mamba pattern—performs best downstream.MVA shares B and C across input channels.Source: §7.2 · pp. 24–25

7.3 Other SSD Extensions from Linear Attention

• Attention-inspired SSD extensions add a configurable kernel feature map, defaulting to SiLU, and optional denominator normalization, but ablations found no significant gain warranting adoption as defaults.The denominator is computed by augmenting X with a column of ones, and the feature map must be positive for the sum to be positive.Source: §7.3 · p. 25

8.1 Tensor Parallel

• Mamba-2’s direct Δ, B, C projections and TP-divisible GroupNorm enable one end-of-block all-reduce per layer, whereas Mamba requires two.Mamba’s Δ, B, C depend on x_c, requiring an intermediate all-reduce before SSM computation; Mamba-2 derives them directly from u.Source: §8.1 · pp. 26–27

8.2 Sequence Parallelism

• Sequence parallelism applies unchanged to Mamba-2 residual and normalization operations; for token mixing, SSM context parallelism passes worker states downstream with communication linear in worker count.Unlike attention’s query-key block interactions, each SSM worker processes its inputs from an initial state and passes its final state to the next worker.Source: §8.2 · pp. 27–28

8.3 Variable Length

• For variable-length batches, Mamba can treat sequences as one long sequence and set A_t = 0 at sequence boundaries to prevent information from crossing between inputs.Source: §8.3 · p. 28

9.2.1 Scaling Laws

• Across ≈125M–≈1.3B parameters, Mamba-2 matches or exceeds Mamba and Transformer++, and Pareto-dominates the Transformer baseline on perplexity, theoretical FLOPs, and wall-clock time.Scaling-law experiments follow the standard Chinchilla protocol.Source: Figure 9 · p. 29, §9.2.1 · p. 30

9.2.2 Downstream Evaluations

• On zero-shot evaluations, Mamba-2 outperforms Mamba at each size and generally matches Pythia at twice its size.The comparisons include open-source models trained for up to 300B tokens; Pile comparisons match the dataset and tokenizer.Source: Table 1 · p. 29, §9.2.2 · p. 30

9.1 Synthetics: Associative Recall

• On MQAR, Mamba-2 performs well across settings while Mamba-1 struggles, and increasing state size from N = 16 to N = 64 or N = 256 consistently improves performance.Mamba-2 also outperforms Mamba-1 when both use N = 16.Source: §9.1 · p. 30

9.2.3 Hybrid Models: Combining SSD Layer with MLP and Attention

• Hybrid SSD-attention models perform best with around 10% attention layers, and six attention layers improve 2.7B model quality over pure Mamba-2 and Transformer++.Adding MLP layers reduces quality but can speed training and inference and facilitate conversion to MoE models.Source: §9.2.3 · pp. 30–31, Table 2 · p. 31, Table 3 · p. 31

9.3 Speed Benchmarks

• SSD runs 2–8× faster than Mamba’s fused scan and beats FlashAttention-2 from sequence length 2K, though whole-model training may be less efficient than Transformer at short lengths.The SSD speedup comes from using matrix multiplication units; Mamba-2 can combine SSD and MLP layers to improve short-sequence training speed.Source: Figure 10 · p. 29, §9.3 · p. 31

9.4.1 Block Design

• Parallel projections save parameters, slightly outperform sequential projections, and enable tensor parallelism; an extra normalization layer slightly improves performance and helps training stability at larger scales.Source: §9.4.1 · p. 32

9.4.2 Head Structure

• Multi-value head patterns differ substantially from multi-query and multi-key patterns despite equal total state size, while parameter-matched multi-head variants perform similarly between them.Source: §9.4.2 · p. 32, Table 5 · p. 33

9.4.3 Attention Kernel Approximations

• The tested kernel approximations do not improve over pointwise nonlinear activations, and adding a normalization term causes instability in most variants.The normalization term slightly improves ReLU; Mamba-2 uses Swish by default, and removing the activation entirely is suggested but was not extensively tested. The negative results may reflect that SSD includes a 1-semiseparable mask absent from the linear-attention methods tested.Source: §9.4.3 · pp. 32–33

10.1 State Space Models

• SSD is a selective SISO SSM with scalar-identity transitions; its restricted transition structure trades some expressivity versus general diagonal A_t for hardware efficiency and easier implementation.The authors hypothesize that refined structured-matrix algorithms could improve efficiency for the general diagonal case as well.Source: §10.1 · p. 34
• Unlike traditional RNNs, SSD’s SISO design enables state expansion to ND, which Mamba found crucial for information-dense language and SSD can increase without slowing the model.Selective SSMs are related to RNN gating mechanisms, including LSTMs and GRUs, but differ in parameterization and state expansion.Source: §10.1 · p. 34

10.2 Structured Matrices

• SSD views SSMs as semiseparable matrix mixers, making their linear-versus-quadratic duality structured versus naive matrix multiplication.The framework’s efficient algorithm uses block decompositions of these matrices; semiseparable-matrix methods and matrix-mixer perspectives may guide further model and algorithm design.Source: §10.2 · pp. 34–35

10.3 (Linear) Attention

• SSD differs from standard attention by omitting softmax for linear scaling and applying an input-dependent 1-semiseparable mask that acts as a replacement and positional signal.The mask’s gates and cumulative products control interactions between positions; structured masked attention generalizes linear attention through different contraction orderings and may enable other efficient variants.Source: §10.3 · p. 35
• SSD does not generalize softmax attention without a finite feature map; its controllable state expansion N compresses history versus attention’s cache scaling with sequence length T.Mamba-2 improves on Mamba on some copying and in-context-learning capabilities, including MQAR, but further understanding is needed.Source: §10.3 · p. 35

10.4 Related Models

• RetNet and TransNormerLLM use decay-based linear attention with parallel, recurrent, and chunkwise algorithms; GateLoop uses input-dependent decay and the same dual quadratic form as SSD.RetNet and TransNormerLLM retain multi-head attention, while Mamba-2 uses the multi-value attention or multi-expand SSM pattern, which the paper reports performs better in Section 9.4.Source: §10.4 · pp. 35–36
• GLA adds data-dependent gates and efficient chunkwise computation; HGRN2 adds state expansion to HGRN’s input-dependent gates, while RWKV-5/6 adopt selectivity and state expansion.Griffin and RecurrentGemma combine input-dependent RNN gating with local attention; Jamba combines Mamba with a few attention layers, and xLSTM adopts state expansion and other gating, normalization, and stabilization techniques.Source: §10.4 · p. 36

11 Conclusion

• The structured state-space duality framework bridges SSMs and attention, and guides Mamba-2 as an architecture at their intersection.The framework also offers insights into SSM performance on language modeling and suggests ways to improve SSMs and potentially Transformers.Source: §11 · p. 36

A Glossary; B.1 Problem Definition; B Efficient Algorithms for the Scalar SSM Scan (1-SS Multiplication)

• The scalar SSM scan is a 1-SS matrix-vector multiplication, and its cumprodsum reduces to cumprod when b=0 or cumsum when a=1.The glossary catalogs frequently used tensor dimensions and matrices or tensors for SSMs and structured masked attention.Source: §B.1 · pp. 42–43, Table 8 · p. 42

B.2 Classical Algorithms; B.2.1 Sequential Recurrence; B.2.2 Parallel Associative Scan

• The sequential recurrence computes the scan one timestep at a time, while the parallel associative scan uses cumulative products of 2 × 2 matrices and divide-and-conquer parallelization.Matrix multiplication is associative, enabling the scan; the appendix also frames efficient sequence-model algorithms as structured matrix multiplication algorithms.Source: B Efficient Algorithms for the Scalar… · p. 42, §B.2.2 · p. 43, §B.2.1 · p. 43

B.3 Efficient Algorithms via Structured Matrix Decompositions

• The structured-decomposition algorithms include dilated, state-passing, fully recurrent, block-decomposition parallel, and recursively reduced equal-block scan modes.They compute the scalar scan through decompositions of the 1-SS matrix, connecting sequence-model algorithms to structured matrix multiplication.Source: §B.3 · p. 44

B.3.1 Dilated Mode

• The dilated mode factors the 1-SS matrix through increasing strides, using O(Tlog(T)) operations but half the depth/span of the work-efficient associative scan.The factorization also shows that 1-SS matrices are a special case of butterfly matrices (Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020).Source: §B.3.1 · p. 44

B.3.2 State-Passing (Chunkwise) Mode; B.3.3 Fully Recurrent Mode

• State-passing computes chunks of arbitrary length and passes each chunk’s final state onward; fully recurrent mode is the special case with chunk size k=1.With linear-time subproblems, state-passing takes linear work but remains sequential; it is useful in general N-semiseparable or batched settings where hidden states can be skipped.Source: §B.3.2 · pp. 44–46, §B.3.3 · p. 46

B.3.4 (Parallel) Block Decomposition Mode

• The block decomposition lowers algorithm depth/span from linear to log(T), trading constant-work combination for O(Tlog(T)) total work.It recursively computes the two self-similar 1-SS subproblems in parallel, then combines them using the rank-1 bottom-left block; recursion may switch to dual quadratic attention on small subproblems.Source: §B.3.4 · p. 46

B.3.5 Associative Scan Mode

• Stage 1 computes each diagonal block’s multiplication results, producing two values while leaving the first element unchanged.For the second block, these values are b2 and a3b2 + b3.Source: §B.3.5 · p. 47
• Stage 2 factors the strictly lower-triangular 2 × 2 blocks as rank-1 matrices and uses Stage 1’s second outputs to form a smaller 1-SS matrix multiplication.The input array b′ has half the size of b.Source: §B.3.5 · p. 47
• Stage 3 broadcasts each Stage 2 answer into two final answers using the left-side column vectors.The passage gives [a2:2 a3:2]ᵀ, [a4:4 a5:4]ᵀ, and [a6:6 a7:6]ᵀ as examples.Source: §B.3.5 · p. 47
• O(T) total work and O(log T) depth/span follow because Stages 1 and 3 take O(T) work while Stage 2 recurses on a half-size self-similar problem.This is the complexity of the three-stage algorithm.Source: §B.3.5 · p. 48
• The matrix-decomposition perspective yields the same computation graph as associative scan without first recognizing a recurrence and defining an associative binary operator.The familiar scan algorithm is described in Appendix B.2.2.Source: §B.3.5 · pp. 47–48
• The algorithm’s computation graph is identical to that of the associative scan described in Appendix B.2.2, though it is derived through structured matrix decomposition.The alternative is to find a structured matrix decomposition for M directly.Source: §B.3.5 · p. 48

C.1 Extras: Closure Properties of SSMs

• Adding an N-SS and a P-SS matrix yields at most an (N+P)-SS matrix; their product is (N+P)-SS, and the inverse of an N-SS matrix is at most (N+1)-SS.The inverse property can be proved using the Woodbury inversion identity, which has featured in structured SSM literature (Gu, Goel, and Ré 2022).Source: §C.1 · p. 48
• Summing two parallel SSMs remains an SSM, while sequentially composing two SSMs yields an SSM whose total state size is additive.The composition result is described as a somewhat nontrivial consequence of the multiplication property.Source: §C.1 · p. 48
• Banded matrices are semiseparable, so their inverses are semiseparable; the fast recurrence properties of semiseparable matrices follow from their inverses being banded.In particular, the inverse of a 1-SS matrix is 2-banded; the semiseparable structure is often motivated by inverting banded matrices.Source: §C.1 · pp. 48–49
• Under certain assumptions, autoregressive structured attention must be an SSM, allowing more general efficient sequence-model families, including attention variants, to be reduced to SSMs.Source: §C.1 · p. 49

C.2 Autoregressive Masked Attention is Semiseparable-Structured Attention

• Efficient autoregressive generation requires constant-time incremental output computation, and this appendix characterizes which structured masked attention instances have that property.The characterization uses a narrow definition of autoregression based on dependence on a fixed number of prior outputs.Source: §C.2 · p. 49
• Order-k autoregression means each output depends only on the current input and the last k outputs; the cumsum matrix is the k=1 case.The proof connects this recurrence to a banded matrix whose inverse is semiseparable.Source: §C.2 · pp. 49–50
• An efficient autoregressive transformation of order k is a state space model of order k+1, so efficient autoregressive attention is semiseparable structured attention.Causality in structured masked attention corresponds to a lower-triangular mask.Source: §C.2 · pp. 49–50

D.1 MQAR Details

• The harder MQAR variant replaces non-query/key/value tokens with random tokens and uses more key-value pairs, longer sequences, and smaller models than the usual variant.It uses sequence lengths T∈{256, 512, 1024}, T/4 key-value pairs, and a total vocabulary size of 8192.Source: §D.1 · p. 50
• MQAR training cycles through datasets with T/32, T/16, T/8, and T/4 key-value pairs, while methods use two layers and sweep model dimensions and learning rates.Training uses eight epochs through each dataset; the attention baseline additionally receives positional embeddings, and learning rates decay linearly each epoch.Source: §D.1 · p. 50

D.2 Scaling Law Details

• Scaling-law models use the GPT2 tokenizer and model sizes specified in Table 9, with training tokens adjusted roughly in proportion to model size.The 1.3B model’s batch size changes from 1M to 0.5M tokens for uniformity, following GPT3 model-size conventions and roughly matching Chinchilla scaling laws.Source: §D.2 · p. 50, Table 9 · p. 51
• The basic scaling-law training recipe uses AdamW, gradient clipping at 1.0, and linear learning-rate warmup with cosine decay.The peak learning rate follows the GPT3 specification by default.Source: §D.2 · pp. 50–51
• The improved recipe uses cosine decay to 1e−5 with a peak learning rate 5× the GPT3 value, no linear biases, RMSNorm, and AdamW β=(.9, .95).The β setting replaces the PyTorch default β=(.9, .999).Source: §D.2 · p. 51

• Mamba-2 is trained on 300B Pile tokens with the GPT-NeoX tokenizer for downstream evaluation.Source: §D.3 · p. 51
• The 1.3B and 2.7B models use batch size 1M; the 2.7B model also follows GPT3 specifications (32 layers, dimension 2560).Otherwise, the models use the same hyperparameters as the scaling experiments.Source: §D.3 · p. 51
• Evaluation uses EleutherAI’s LM evaluation harness on the same tasks as Mamba, with one additional task: LAMBADA (Paperno et al. 2016), HellaSwag, PIQA (Bisk et al. 2020), and ARC-challenge (P. Clark et al. 2018).Source: §D.3 · p. 51
• The evaluation suite also includes ARC-easy, an easy subset of ARC-challenge, WinoGrande, and OpenBookQA (Mihaylov et al. 2018).Source: §D.3 · p. 51

D.4 Ablation Details

• Based approximates the exp kernel with the quadratic Taylor expansion exp(x) ≈ 1 + x + x^2/2.The paper includes Based (Arora, Eyuboglu, Zhang, et al. 2024) in its ablations.Source: §D.4 · p. 51
• ReBased uses ψQuadratic(x) = x⊗x, corresponding to the x^2 kernel transformation, and applies layer normalization beforehand; the ablations compare combinations with Swish.The paper treats layer normalization as an alternative nonlinear activation to its default Swish activation; ReBased is attributed to Aksenov et al. 2024.Source: §D.4 · p. 51
• For the 130M and 380M models, the ablations use N = 64 and N = 256 respectively, project to 8 and 16 dimensions before the (Re)Based feature maps, and adjust layer counts for fewer parameters.The resulting total state sizes for 130M models are 64 for ReBased and 73 for Based.Source: §D.4 · p. 52