State Space Model vs Attention Complexity

mediumFREEVision MambaAlgorithms, ComplexityAdvanced Architectures

State Space Models (SSMs) like Mamba offer an alternative to Transformers with linear complexity in sequence length.

Self-Attention complexity: $O(N^2 \cdot D)$ - quadratic in sequence length $N$

SSM (Mamba) complexity: $O(N \cdot D \cdot S)$ - linear in sequence length

$N$ : Sequence length
$D$ : Hidden dimension
$S$ : State dimension (typically 16-64)

Task: Return the asymptotic complexity for both architectures as a tuple of strings.

This linear scaling enables Mamba to process much longer sequences (e.g., 1M tokens) that would be computationally infeasible for Transformers.

Example

Input

N=1000, D=512, S=16

Output

("O(N*D*S)", "O(N^2*D)")

Reasoning

SSM processes each token once with state updates: O(N×D×S). Attention computes all pairwise interactions: O(N²×D). For N=1000, SSM is ~60× more efficient.

Constraints

$N$ (Sequence Length): $1 \leq N \leq 10^6$
$D$ (Hidden Dimension): $1 \leq D \leq 4096$
$S$ (State Dimension): $1 \leq S \leq 64$

More from Advanced Architectures

Complete Mixer Block with Residualshard Channel Mixing with GELUmedium Token Mixing Layermedium Bidirectional Scan for Visionmedium Discretize Continuous State Space Modelhard

Concept primer

State Space Model vs Attention Complexity

Background Knowledge

State Space Models (SSMs) are a family of sequence models inspired by continuous-time dynamical systems from control theory. Unlike Transformers which use self-attention with $O(N^2)$ complexity, SSMs process sequences with linear complexity $O(N)$ .

The Complexity Problem with Attention

Self-attention computes pairwise interactions between all tokens:

Attention(Q, K, V) = softmax(QK^T / √d) × V

For sequence length $N$ and hidden dimension $D$ :

Memory: $O(N^2)$ to store attention matrix
Compute: $O(N^2 \cdot D)$ for matrix multiplications

How SSMs Achieve Linear Complexity

SSMs model sequences as discretized continuous systems:

h(t) = Ah(t-1) + Bx(t) # State update
y(t) = Ch(t) # Output projection

With state dimension $S$ :

Memory: $O(N \cdot S)$ - only store current state
Compute: $O(N \cdot D \cdot S)$ - linear in sequence length

Algorithm/Approach

To compare complexities:

Attention: $O(N^2 \cdot D)$ operations
SSM: $O(N \cdot D \cdot S)$ operations

When $S < N$ (state size smaller than sequence), SSMs are more efficient.

Common Pitfalls

Confusing state dimension $S$ with hidden dimension $D$
Forgetting that attention complexity is quadratic in $N$ , not $D$
Not accounting for the constant factors in Big-O notation

Time & Space Complexity

Model	Time	Space
Attention	O(N²D)	O(N² + ND)
SSM	O(NDS)	O(NS + DS)

solution.py

Python 3.13

Test results