Skip to content

State Space Models (SSMs) Technical Notes

Quick Reference

  • Definition: State Space Models (SSMs) are continuous-time linear dynamical systems parameterized for deep learning, enabling efficient sequence modeling through structured approximations and parallelizable computations.
  • Key Use Cases: Foundation models for long-context multimodal data, continuous-time forecasting with irregular sampling, and hardware-efficient alternatives to transformers in large-scale training.
  • Prerequisites: Advanced linear algebra, differential equations, proficiency in PyTorch/JAX for custom kernels, and experience with sequence models like transformers.

Table of Contents

  1. Introduction
  2. Core Concepts
  3. Implementation Details
  4. Real-World Applications
  5. Tools & Resources
  6. References
  7. Appendix

Introduction

What

State Space Models (SSMs) represent sequences as discretized continuous-time systems, with advanced variants like S4, Mamba, and Hyena using structured parameterizations for O(N) time/space complexity in training and inference.

Why

SSMs overcome transformer's quadratic scaling, enabling modeling of million-length sequences with constant memory, while supporting continuous-time dynamics and selective context integration.

Where

SSMs power next-generation foundation models in NLP (long-document QA), genomics (DNA sequences), and control (continuous robotics), often in hybrid architectures with attention.

Core Concepts

Fundamental Understanding

  • Basic Principles: SSMs model dx/dt = Ax + Bu, y = Cx + Du, with discretization to x_{k+1} = \bar{A} x_k + \bar{B} u_k, y_k = C x_k; advanced forms use HiPPO initialization and DPLR for long-range modeling.
  • Key Components:
  • Structured Parameterizations: Diagonal plus low-rank (DPLR) for A to enable fast matrix powers via Cauchy kernels.
  • Selective Mechanisms: Input-dependent gating (e.g., S6) for context selection, akin to attention but O(N).
  • Discretization Schemes: Learnable dt or adaptive methods for handling variable sampling rates.
  • Hybrid Architectures: Stacking SSM layers with MLPs or attention for enhanced expressivity.
  • Common Misconceptions:
  • Linearity limits expressivity: Non-linear activations and gating enable universal approximation.
  • Fixed kernels suffice: Learnable structured A matrices adapt to data.
  • Inference-only efficiency: FFT/parallel scans enable fast training.

Visual Architecture

graph TD
    A["Continuous Dynamics<br>dx/dt = Ax + Bu"] -->|"HiPPO Init & DPLR"| B[Structured Params]
    B -->|Learnable Discretization| C["Discrete \bar{A}, \bar{B}, C"]
    C -->|Selective Gating| D[Input-Dependent Scan]
    D -->|Parallel Computation| E[Output y_k Sequence]
    E -->|"Stacking & Hybrids"| F[Deep SSM Model]
- System Overview: Continuous dynamics are structured and discretized, selectively scanned in parallel, and stacked into deep models for complex sequence tasks. - Component Relationships: Structuring enables efficient discretization, gating adds selectivity, parallel scans scale computation.

Implementation Details

Advanced Topics

import torch
import torch.nn as nn
from einops import rearrange, repeat

class SSM(nn.Module):
    def __init__(self, dim, state_dim, heads=8):
        super().__init__()
        self.heads = heads
        self.dim = dim // heads

        # Structured params (DPLR)
        self.Lambda = nn.Parameter(torch.randn(heads, state_dim) - 0.5)
        self.P = nn.Parameter(torch.randn(heads, state_dim))
        self.Q = self.P  # Conjugate symmetry

        self.B = nn.Parameter(torch.randn(heads, state_dim, self.dim))
        self.C = nn.Parameter(torch.randn(heads, self.dim, state_dim))

        # Selective gating
        self.Dt = nn.Linear(dim, heads)  # Learnable timestep
        self.Gate = nn.Linear(dim, heads * state_dim)  # Input-dependent

    def discretize(self, dt):
        # DPLR to discrete \bar{A}, \bar{B}
        A = torch.diag_embed(self.Lambda) - torch.einsum('h m, h n -> h m n', self.P, self.Q)
        A_bar = torch.exp(A * dt.unsqueeze(-1).unsqueeze(-1))
        B_bar = (torch.eye(A.size(-1)) - A_bar) @ self.B * dt.unsqueeze(-1)
        return A_bar, B_bar

    def forward(self, u):
        # u: (batch, seq, dim)
        u = rearrange(u, 'b l (h d) -> b h l d', h=self.heads)

        dt = nn.functional.softplus(self.Dt(u.mean(dim=2)))  # (b h)
        gate = self.Gate(u).sigmoid()  # (b l h m) -> rearrange if needed

        A_bar, B_bar = self.discretize(dt)

        # Selective parallel scan
        states = torch.zeros(u.size(0), self.heads, u.size(2), A_bar.size(-1), device=u.device)
        for t in range(u.size(2)):  # Parallelize in prod with cumprod
            states[:, :, t] = A_bar @ states[:, :, t-1] + B_bar @ u[:, :, t].unsqueeze(-1)
            states[:, :, t] *= gate[:, t].unsqueeze(-1)  # Selective forget

        y = torch.einsum('b h l m, h d m -> b l h d', states, self.C)
        y = rearrange(y, 'b l h d -> b l (h d)')
        return y

# Usage in deep model
class DeepSSM(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.ssm = SSM(dim, state_dim=64)
        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim))

    def forward(self, x):
        return self.mlp(self.ssm(x)) + x  # Residual
- System Design: - Multi-Head SSMs: Parallel heads for increased capacity. - Hybrid Layers: Combine with MLPs for non-linearity. - Custom Kernels: Use Triton/JAX for fast parallel scans. - Optimization Techniques: - HiPPO Initialization: For long-range dependencies. - DPLR Approximation: Enables closed-form matrix powers. - Gating for Selectivity: Input-dependent forgetting rates. - Production Considerations: - Scaling Laws: Train on massive datasets for emergent abilities. - Hardware Efficiency: Optimize for TPUs with XLA compilation. - Continuous Mode: Handle variable dt for irregular data.

Real-World Applications

Industry Examples

  • Use Case: Genomic sequence modeling with million-base-pair contexts.
  • Implementation Pattern: Multi-head selective SSMs with hybrid layers.
  • Success Metrics: State-of-the-art perplexity on long sequences.

Hands-On Project

  • Project Goals: Build a scalable SSM for long-context NLP.
  • Implementation Steps:
  • Implement DPLR-structured SSM with selective gating.
  • Stack into deep residual model.
  • Train on BookCorpus or similar.
  • Evaluate on long-range tasks.
  • Validation Methods: Measure perplexity scaling with length.

Tools & Resources

Essential Tools

  • Development Environment: JAX or PyTorch with CUDA.
  • Key Frameworks: Flax/Haiku for JAX models.
  • Testing Tools: Hugging Face Datasets.

Learning Resources

  • Documentation: Mamba GitHub repo.
  • Tutorials: "Implementing S4 from Scratch".
  • Community Resources: LessWrong AI alignment forums.

References

  • S4 Paper: Gu et al., 2021.
  • Mamba Paper: Gu & Dao, 2023.
  • Hyena Paper: Poli et al., 2023.

Appendix

Glossary

  • DPLR: Diagonal Plus Low-Rank matrix structure.
  • HiPPO: High-order Polynomial Projection Operators.
  • Selective Scan: Input-dependent state recurrence.

Setup Guides

  • Install JAX: pip install jax jaxlib.