PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tri Dao, Albert Gu

arXiv:2405.21060v1cs.LG

TL;DR

Transformers and structured state-space models are connected through structured-matrix formulations, addressing efficiency differences in sequence modeling. The SSD framework develops these connections and guides Mamba-2, which matches or outperforms Mamba and open-source Transformers on standard downstream evaluations.

  • Problem

    Transformers face quadratic training costs and sequence-length cache growth, while SSMs offer linear training scaling and constant generation state size.

  • Method

    SSD represents sequence models through structured matrices and tensor contractions, establishing connections between SSMs, semiseparable matrices, and attention variants.

  • Results

    Mamba-2 matches or outperforms Mamba and open-source Transformers on standard downstream evaluations, including a 2.7B model outperforming several larger or similarly sized baselines.

  • Takeaways & Limitations

    SSD provides theoretical and algorithmic tools for transferring attention-derived design and systems ideas to SSMs, guiding Mamba-2 at the intersection of SSMs and structured attention.

  • Takeaways & Limitations

    SSD does not generalize standard softmax attention or attention-kernel transformations without a finite feature map, and tradeoffs in capabilities remain incompletely understood.

Abstract

from arXiv · show

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.

1 Introduction

The paper develops structured state space duality (SSD) to connect SSMs with attention through structured matrices, then uses the framework to design faster, competitive Mamba-2 models.

  • Structured State Space Duality: SSMs and attention are connected through structured matrices, enabling theoretical and algorithmic transfers between the two model families.The framework represents sequence models through matrix transformations and tensor contractions.
  • Theoretical Connections: The framework identifies semiseparable matrices as the central bridge between state space models and attention variants.It also relates SSMs to structured masked attention and establishes an intersection with linear attention.
  • Efficient Algorithms: SSD introduces efficient algorithms that combine the linear SSM recurrence with a quadratic dual form and exploit modern matrix-multiplication hardware.The dedicated SSD implementation is described as 2–8× faster than optimized selective-scan implementations.
  • Architecture Design: Mamba is interpreted as a multi-input SSM analogous to multi-value attention, motivating attention-inspired head structures and parallel block modifications.The modifications support tensor parallelism and move data-dependent projections to the beginning of the block.
  • Empirical Results: Mamba-2 Pareto dominates Mamba and Transformer++ in perplexity and wall-clock time under the investigated Chinchilla scaling setting.On the Pile, a 2.7B-parameter Mamba-2 trained on 300B tokens outperforms several larger or comparable baseline models on standard downstream evaluations.

2 Background and Overview

The paper situates structured SSMs among recurrent, convolutional, and continuous-time sequence models, then uses structured-matrix ideas to motivate SSD and Mamba-2’s efficiency trade-offs.

  • SSM Background: Structured SSMs map sequences through latent states and can be viewed as recurrent neural networks with linear input recurrences.Their state size N controls the model’s expansion and information capacity.
  • SSM Variants: Time-invariant SSMs are equivalent to convolutions, with globally parameterized kernels, and can switch between convolutional training and recurrent autoregressive inference.Selective SSMs allow time-varying parameters that choose which inputs to emphasize or ignore, but require hardware-aware implementations.
  • Structured Matrix View: SSD extends the matrix perspective by treating SSMs as structured transformations and generalizing linear attention into structured masked attention.The framework uses structured matrices to connect SSM-like linear computation with attention-like forms.
  • SSD Algorithm: The proposed SSD algorithm uses block decompositions of the SSM transformation matrix to improve efficiency trade-offs beyond purely linear or quadratic computation.The method is designed to be hardware-efficient and simpler to implement than general selective SSMs.

3 State Space Models are Structured Matrices

This section proves that SSM transformations are semiseparable matrix multiplications, yielding linear-time algorithms and an explicit matrix-based view of SSM computation.

  • Semiseparable Structure: An N-semiseparable matrix is lower triangular with every submatrix on or below the diagonal having rank at most N.The order N is the semiseparable matrix’s rank bound.
  • Representations: Semiseparable matrices and sequentially semiseparable representations are equivalent, and the representation can be compressed from O(N^2T) parameters to O(NT).The paper uses this representation throughout its algorithmic development.
  • Scalar Recurrence: 1-SS matrices correspond to scalar SSM recurrences, whose multiplication can be computed using cumulative product-sum operations.These scalar recurrences serve as building blocks for the paper’s main algorithms.
  • SSMs as Semiseparable Matrices: An SSM with state size N is identical to multiplication by an N-SS matrix in sequentially semiseparable representation.This equivalence is the section’s central structural result.
  • Efficiency Results: An N-SS matrix of size T uses O(NT) parameters and supports matrix-vector multiplication in O(NT) time and space.The result supplies the efficiency basis for computing SSM transformations.
  • Efficiency Results: Even unstructured SSM dynamics can be computed in O(TN) time after preprocessing, matching the lower bound implied by the sizes of B and C.The result applies beyond diagonalizable SSMs.

4 Structured Masked Attention: Generalizing Linear Attention with Structured Matrices

Structured masked attention generalizes linear attention by applying kernel attention through structured masks, preserving efficient computation when the mask supports fast matrix multiplication. A tensor-contraction view explains linear attention’s recurrence and exposes a broader duality between quadratic and subquadratic forms.

  • 4.3 Structured Masked Attention: Structured masked attention applies a mask matrix L to kernel attention, producing M = QK^⊤◦L and the sequence transformation Y = MV.
  • 4.3 Structured Masked Attention: Structured masked attention extends prior kernel-attention variants, including Linear Attention and RetNet, while permitting further structured-matrix instantiations.
  • 4.2 A Tensor Contraction Proof of Linear Attention: Autoregressive kernel attention can be computed in O(T) time by a recurrence with constant time per step.
  • 4.2 A Tensor Contraction Proof of Linear Attention: The linear-time recurrence follows because multiplication by the causal lower-triangular mask is a feature-wise cumulative sum.
  • 4.3 Structured Masked Attention: Any structured mask with subquadratic matrix-vector multiplication can preserve linear-attention efficiency by accelerating the bottleneck involving both sequence axes.
  • 4.3 Structured Masked Attention: Different contraction orderings yield quadratic and subquadratic forms for linear attention, reflecting the same underlying duality as state-space models.

5 State Space Duality

Structured state-space models and structured masked attention admit dual linear and quadratic formulations. Their key shared case is scalar-identity SSMs paired with 1-semiseparable masks, while general semiseparable attention is more expressive than standard SSMs.

  • Scalar structured SSMs: The scalar structured SSM’s naive quadratic computation is exactly equivalent to quadratic masked kernel attention.This equivalence follows by materializing the semiseparable matrix and performing quadratic matrix-vector multiplication.
  • The SSD correspondence: 1-semiseparable SMA is a special case of a diagonal SSM whose diagonal matrix is a scalar multiple of the identity.The shared structure removes the extra state dimension present in the general diagonal SSM formulation.
  • Autoregressive attention: Efficient autoregressive structured masked attention must use a semiseparable mask, making efficient autoregressive attention general semiseparable SMA.This converse result characterizes the masks compatible with bounded-order autoregression.
  • Scope of the correspondence: General semiseparable SMA is strictly more expressive than 1-semiseparable SMA and cannot be described by a standard SSM.Its multiplication and linear-form expansions can nevertheless be absorbed into a larger 1-SS SMA instance.
  • Dual computational forms: SSMs have both a linear recurrent form and a quadratic matrix form, while attention variants have both quadratic pairwise and linear tensor-contraction forms.These forms arise from different representations or computation orders.
  • The SSD correspondence: Scalar-identity SSMs and 1-semiseparable structured masked attention are duals with exactly the same linear and quadratic forms.This is the central correspondence summarized by SSD.

6 A Hardware-Efficient Algorithm for SSD Models

The SSD algorithm computes state space models by combining recurrent and quadratic modes through a block decomposition of semiseparable matrices. This yields linear sequence-length scaling, matrix-multiplication-friendly computation, and tight resource bounds.

  • Block decomposition: The SSD algorithm combines linear recurrent computation with quadratic attention-like computation through block decomposition of semiseparable matrices.Diagonal blocks use the quadratic dual form, while off-diagonal blocks use low-rank factorizations through the hidden state.
  • Computational guarantees: O(TN^2) training FLOPs, O(TN) inference FLOPs, and O(N^2) inference memory are achieved with work dominated by matrix multiplications.These bounds are tight for state expansion N and head dimension P = N.
  • Block decomposition: Diagonal blocks are independent smaller SSM problems that can use quadratic SMA computation and execute in parallel.Off-diagonal blocks are low-rank by the defining property of semiseparable matrices.
  • Block decomposition: The block decomposition separates each chunk’s output into intra-chunk effects from current inputs and inter-chunk effects from prior inputs.The inter-chunk contribution is factored through the SSM hidden state.
  • Computational cost: The resulting computation uses O(TN^2) FLOPs and O(TN) memory, while its scalar SSM scan is parallelizable and negligible in the GPU implementation.Blocking reduces the scan length from T to T/Q, making it Q times cheaper than a pure SSM scan.

7 The Mamba-2 Architecture

Mamba-2 applies the SSD framework and attention-inspired design choices to simplify the Mamba block and organize its sequence-transformation heads. The resulting architecture uses parallel parameter projections, added normalization, and a multi-value attention-style head pattern.

  • Block design: Mamba-2 produces A, B, C, and X in parallel at the beginning of the block, analogous to parallel Q, K, and V projections in attention.This replaces Mamba-1’s sequential data-dependent projections.
  • Block design: An additional normalization layer is placed before the final output projection to alleviate instabilities observed in larger models.The change follows the placement used in NormFormer and related linear-attention models.
  • Head patterns: The state size N and head dimension P correspond to attention’s QK and V head dimensions, respectively, and are generally kept near 64 or 128.Increasing model dimension D therefore increases the number of heads while keeping N and P fixed.
  • Head patterns: The Mamba-2 block uses a multi-value attention or multi-input SSM pattern, sharing B and C across the input channels X.This pattern treats X as the main SSM input and B and C as shared parameters.
  • Head patterns: Among controlled head-pattern variants, the MVA pattern used in Mamba performs best in downstream evaluations.The comparison controls parameter counts and total state dimension.
  • Kernel features: Linear-attention-inspired architectural modifications did not improve performance enough to become default settings, so the kernel feature map remains a hyperparameter.The paper evaluates these choices as a negative result in its ablations.

8 Systems Optimization for SSMs

The systems design adapts Mamba-2 to large-scale training through tensor and sequence parallelism. Parallel projections reduce tensor-parallel communication, while SSM state passing enables sequence partitioning with communication linear in the number of workers.

  • Tensor parallelism: Mamba-1 tensor parallelism requires two all-reduces because Δ, B, and C depend on the complete projected input, doubling communication relative to Transformers.Communication can already occupy 10–20% of large-scale Transformer training time.
  • Tensor parallelism: Mamba-2’s parallel projections allow tensor parallelism with one all-reduce per block, matching the communication pattern of Transformer attention and MLP blocks.The projections for Δ, B, and C are obtained directly from u, allowing them to be split across devices.
  • Tensor parallelism: GroupNorm is arranged so that tensor-parallel groups need no communication within the block for the normalization operation.The number of groups is chosen to be divisible by the tensor-parallel degree.
  • Sequence parallelism: Sequence parallelism also applies to Mamba-2 residual and normalization operations through the same reduce-scatter and all-gather decomposition used for Transformer training.The architecture’s matching residual and normalization structure allows SP to apply without modification.
  • Sequence parallelism: SSM context parallelism splits a sequence across workers by passing each worker’s final state to the next, giving communication bandwidth linear in the number of workers.This is the sequence analogue of the SSD block decomposition.

9 Empirical Validation

The empirical evaluation tests Mamba-2 on recall, language modeling, hybrid architectures, and efficiency. Across these settings, Mamba-2 improves over Mamba, benefits from attention hybrids, and its SSD algorithm delivers substantial speedups.

  • Language Modeling: Mamba-2 Pareto dominates Mamba and Transformer++ in perplexity and wall-clock time under Chinchilla scaling laws.The scaling-law experiments use models from approximately 125M to 1.3B parameters trained on the Pile.
  • Language Modeling: Mamba-2 outperforms Mamba in zero-shot evaluations and generally matches Pythia models at twice its model size.The comparison covers open-source language models trained for up to 300B tokens.
  • Synthetic Associative Recall: Mamba-2 performs well across all MQAR settings, whereas Mamba-1 struggles, including when state size is controlled at N = 16.Increasing state size from N = 16 to N = 64 and N = 256 consistently improves MQAR performance.
  • Hybrid Models: Adding a few attention layers improves quality over pure Mamba-2 and Transformer++, with around 10% attention layers performing best.The authors report that combining SSD, attention, and MLP layers outperforms either pure Transformer++ or Mamba-2.
  • Efficiency Benchmarks: 2-8× faster than Mamba’s fused scan, SSD also becomes faster than FlashAttention-2 at sequence length 2K and above.SSD gains speed by reformulating computation to use GPU matrix-multiplication units.
  • Efficiency Benchmarks: Mamba-2 may be less efficient to train than Transformer at short sequence lengths such as 2K.The authors attribute this comparison partly to the hardware efficiency of Transformer MLP layers.

10 Related Work and Discussion

The discussion places SSD at the intersection of SSMs, structured matrices, and attention. It emphasizes both the framework’s efficiency and its boundaries relative to general softmax attention and broader sequence-model designs.

  • SSM Design Space: SSD is a time-varying selective SSM with SISO dimensions and scalar-identity transition structure.Its restriction on transition expressivity improves hardware friendliness compared with more general diagonal transitions.
  • Structured-Matrix View: SSD characterizes SSMs as semiseparable matrix mixers, reframing their linear and quadratic forms as structured and naive matrix multiplication.This perspective connects efficient sequence transformations to structured matrices with sub-quadratic complexity.
  • SSD and Attention: Compared with causal attention, SSD omits softmax and uses an input-dependent 1-semiseparable mask that replaces it.The mask’s cumulative products act as gates controlling interactions between sequence positions.
  • Structured Masked Attention: Structured masked attention generalizes linear attention through alternative contraction orderings and supports other efficient attention variants.The framework is broader than SSD because different structured masks can yield different properties.
  • Scope and Limitations: SSD does not generalize standard softmax attention and instead uses a controllable state expansion factor to compress history.The paper notes that Mamba-2 improves some copying and in-context-learning capabilities, while more remains to be understood.
  • Related Sequence Models: The discussion identifies related models such as RetNet and TransNormerLLM as SSD instances with time-invariant transitions.In the structured-mask view, their masks correspond to decay matrices.

11 Conclusion

The paper concludes that SSD bridges SSMs, structured matrices, and attention, and uses these connections to guide Mamba-2. Its algorithmic perspective derives multiple equivalent ways to compute SSM scans and supports parallel implementations.

  • Conclusion: SSD bridges SSMs, structured matrices, and attention, providing a framework for understanding and improving sequence models.The framework directly guides the design of Mamba-2 at the intersection of selective SSMs and structured attention.
  • Scalar SSM Scan: The scalar SSM scan can be viewed equivalently as a recurrent scan, a 1-SS matrix-vector multiplication, or a cumulative product-sum.The cumprodsum reduces to cumulative product when b = 0 and cumulative sum when a = 1.
  • Parallel Scan Algorithms: Associative scans parallelize SSM recurrences because the recurrence is represented using associative matrix multiplication.Divide-and-conquer algorithms exploit this associativity.
  • Computation Modes: The paper presents recurrent, state-passing, dilated, block-decomposition, and scan computation modes for SSMs.These modes arise from structured matrix decompositions of the 1-SS matrix.
  • Structured Decompositions: The 1-SS matrix factorization is closely related to dilated convolutions and makes 1-SS matrices a special case of butterfly matrices.These connections relate SSM computation to other structured-matrix algorithms.
  • Parallel Algorithms: Block decomposition lowers algorithmic depth from linear to log(T) by recursing on both subproblems in parallel.The tradeoff is linear rather than constant work in the combination step.

B.3.5 Associative Scan Mode

The associative scan mode derives a parallel state-space computation through a three-stage structured matrix factorization. Its stages preserve linear total work while reducing span to logarithmic depth.

  • Complexity: The state-passing chunkwise algorithm has linear work but still includes sequential operations, motivating the parallel scan formulation.Other parallelizable block-reduction and dilated-mode methods achieve logarithmic span but perform O(Tlog(T)) work.
  • Three-stage factorization: The algorithm proceeds in three stages: diagonal-block computation, recursive reduction of lower-triangular blocks, and broadcasting through left-side column vectors.Stage 1 computes diagonal-block answers; Stage 2 forms a half-size self-similar problem; Stage 3 broadcasts intermediate answers.
  • Three-stage factorization: Stage 2 factors each 2 × 2 block in the strictly lower-triangular region as a rank-1 matrix and reduces the problem recursively.The right-side row vectors match bottom rows from diagonal blocks, so Stage 1 supplies part of the required intermediate results.
  • Complexity: Stage 1 and Stage 3 require O(T) work, while Stage 2 reduces to a self-similar problem of half the size.This recursive structure yields O(T) total work and O(logT) depth/span.
  • Interpretation: The factorization is equivalent to the associative scan computation graph, offering a matrix-decomposition perspective on the same parallel algorithm.The matrix view replaces recognizing a recurrence and deriving an associative operator with directly finding a structured matrix decomposition.

C.1 Extras: Closure Properties of SSMs

Semiseparable closure properties transfer directly to state-space models, explaining how SSMs behave under addition, composition, and inversion. These properties also connect SSMs with efficient attention variants.

  • Matrix closure: Semiseparable matrices are closed under addition, multiplication, and inversion, with orders bounded by N+P, N+P, and N+1 respectively.The inverse bound can be obtained using the Woodbury inversion identity.
  • SSM consequences: These matrix closure results imply corresponding closure properties for state-space models.The paper uses the semiseparable representation to transfer algebraic structure from matrices to SSMs.
  • SSM consequences: Summing two parallel SSMs remains an SSM, while sequentially composing two SSMs remains an SSM with additive total state size.The additive state-size result is described as a nontrivial consequence of multiplication closure.
  • SSM consequences: Because banded matrices are semiseparable and their inverses are semiseparable, fast recurrence properties can be understood through banded inverses.The paper notes that semiseparable structure is often motivated by taking inverses of banded matrices.
  • Attention connection: Under the stated conditions, autoregressive structured attention can be represented as an SSM, extending the reduction to efficient attention variants.The paper explicitly relates autoregressive structured attention to semiseparable SSMs.

C.2 Autoregressive Masked Attention is Semiseparable-Structured Attention

The paper characterizes efficient autoregressive masked attention through semiseparable structure. Under an order-k autoregressive dependence, the transformation is an SSM of order k+1.

  • Problem setup: The section asks which structured masked-attention instances support efficient autoregression.The analysis uses a narrow autoregressive definition aligned with classical time-series formulations.
  • Problem setup: Causality corresponds to the attention mask being lower triangular, while order-k autoregression restricts each output to the current input and previous k outputs.This defines the local dependence pattern studied in the characterization.
  • Example: The cumulative-sum mask is a special case with k=1, yielding y_t = x_t + y_t−1.This provides a concrete example within the autoregressive characterization.
  • Main theorem: An efficient autoregressive transformation of order k is a state-space model of order k+1.The theorem formalizes the equivalence between this autoregressive structure and an SSM representation.
  • Proof strategy: The proof converts the autoregressive recurrence into a banded matrix and uses semiseparable closure plus the SSM characterization theorem.A slightly stronger semiseparable bound of k+1 is obtained from additional banded-matrix structure.
  • Conclusion: The paper concludes that efficient autoregressive attention is semiseparable structured matrix attention.The characterization connects the attention formulation to the semiseparable matrix framework.

D.3 Downstream Evaluation Details

The downstream evaluation trains Mamba-2 on the Pile and evaluates it with the LM evaluation harness across the tasks used for Mamba, plus an additional task.

  • Training: Mamba-2 is trained on 300B tokens from the Pile for downstream evaluation.The training uses the GPT-NeoX tokenizer.
  • Model settings: For the 2.7B model, the evaluation follows GPT3 specifications with 32 layers and dimension 2560.The downstream setup otherwise uses the scaling-experiment hyperparameters, with batch size 1M for the 1.3B and 2.7B models.
  • Evaluation suite: Downstream evaluation uses the EleutherAI LM evaluation harness on the Mamba task suite with one additional task.The listed tasks include LAMBADA, HellaSwag, PIQA, ARC-challenge, ARC-easy, WinoGrande, and OpenBookQA.
  • Evaluation suite: The evaluation includes WinoGrande and OpenBookQA alongside the listed language-understanding benchmarks.These two tasks are explicitly included in the downstream task set.

D.4 Ablation Details

The ablations compare Based and ReBased feature-map variants, including layer normalization as an alternative to Swish, alongside zero-shot evaluations showing Mamba-2’s relative performance.

  • (Re)Based Details: Based approximates the exponential kernel with a quadratic Taylor expansion, while ReBased uses a quadratic feature map and layer normalization.The ablations examine these alternatives alongside the default Swish activation.
  • (Re)Based Details: Layer normalization is treated as an alternative non-linear activation to Swish, and combinations of these choices are ablated.
  • Zero-shot Evaluations: Mamba-2 outperforms Mamba at each model size and generally matches Pythia models at twice its size in zero-shot evaluations.The comparison covers open-source language models trained for up to 300B tokens, with Pile comparisons restricted to the same dataset and tokenizer.
Loading 2405.21060v1…