Mamba и State Space Models¶

~10 минут чтения

Предварительно: Реализация внимания с нуля | Эффективные трансформеры

Self-attention квадратичен: $O(n^2 \cdot d)$. Для 1M токенов это 1 триллион пар, а KV cache на 70B модели при 128K съедает 64 GB. State Space Models (SSM) сжимают всю историю в fixed-size state -- $O(n)$ по времени, $O(1)$ по памяти при инференсе: 16 MB вместо 512 GB KV cache. Mamba добавляет к SSM selection mechanism -- параметры зависят от входа, и модель сама решает что запомнить ("Python"), а что забыть ("the"). При 3B параметрах Mamba матчит Transformer по perplexity, но работает в 3--5x быстрее на длинных последовательностях. Тренд 2026: гибриды (Jamba: 7 Mamba + 1 Attention слой, 3x throughput).

Ключевые концепции¶

Проблема: O(n^2) attention¶

Self-attention квадратична: $O(n^2 \cdot d)$. Для 1M tokens = 1T пар. KV cache линейно растет, 128K на 70B модели = 64 GB.

State Space Model (SSM): сжимает всю историю в fixed-size state. O(n) time, O(1) memory при инференсе.

Эволюция¶

2021: S4 (Structured State Space) -- Gu et al.
2022: H3, Hyena -- hybrid approaches
2023: Mamba -- selective SSM (Dec 2023)
2024: Mamba-2 (SSD), RWKV-5/6, Jamba, Griffin
2025: Mamba-3, RWKV-7 (Goose), Nemotron-H, Bamba
2026: Hybrid adoption, Jamba Reasoning

SSM: математика¶

Continuous-time: $$h'(t) = Ah(t) + Bx(t)$$

\[y(t) = Ch(t) + Dx(t)\]

Discrete-time (для последовательностей): $$h_t = \bar{A}h_{t-1} + \bar{B}x_t$$

\[y_t = Ch_t + Dx_t\]

Discretization: $$\bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B$$

Где: - $h_t \in \mathbb{R}^N$ -- hidden state (fixed-size) - $A \in \mathbb{R}^{N \times N}$ -- state transition matrix - $B \in \mathbb{R}^{N \times 1}$ -- input matrix - $C \in \mathbb{R}^{1 \times N}$ -- output matrix - $\Delta$ -- discretization step

HiPPO Matrix (основа long-range memory)¶

\[A_{nk} = -\begin{cases} (2n+1)^{1/2}(2k+1)^{1/2} & \text{if } n > k \\ n+1 & \text{if } n = k \\ 0 & \text{if } n < k \end{cases}\]

Проецирует историю входов на базис ортогональных полиномов -- позволяет реконструировать прошлые сигналы.

S4: Structured State Space¶

Первый практичный SSM (Gu et al., 2021): - Diagonal + low-rank структура для матрицы A - Parallel scan для O(T log T) обучения - Continuous-time -> discrete-time conversion

Ограничение: статические параметры (нет зависимости от входа).

Complexity comparison¶

Операция	Transformer	Mamba (SSM)
Training	O(n^2 * d)	O(n * d^2)
Inference (per token)	O(n * d)	O(d^2)
Memory (training)	O(n^2)	O(n * d)
Memory (inference)	O(n * d) (KV cache)	O(d) (fixed state)

Mamba: Selective State Space Model¶

Ключевая инновация: input-dependent parameters¶

Standard SSM: A, B, C фиксированы -- модель не может быть селективной.

Mamba: параметры зависят от входа: $$\Delta_t = f_\Delta(x_t), \quad B_t = f_B(x_t), \quad C_t = f_C(x_t)$$

\[h_t = \exp(\Delta_t A) \cdot h_{t-1} + \Delta_t B_t \cdot x_t\]

\[y_t = C_t \cdot h_t\]

Модель решает ЧТО запомнить: пропускает "the", "and", запоминает ключевые токены. Dynamic state compression.

Mamba Block¶

graph TD
    IN["Input"] --> LIN1["Linear (proj to d_inner)"]
    LIN1 --> CONV["Conv1d (local context)"]
    CONV --> SSM["Selective SSM"]
    SSM --> GATE["Gating (y * SiLU(res))"]
    GATE --> LIN2["Linear (proj back)"]
    LIN2 --> OUT["Output"]

    SSM --- DETAIL["SiLU activation<br/>x_proj -> (Delta, B, C)<br/>dt_proj -> discretization<br/>h = Ah + Bx"]

    LIN1 -.->|residual| GATE

    style IN fill:#e8eaf6,stroke:#3f51b5
    style SSM fill:#e8f5e9,stroke:#4caf50
    style GATE fill:#fff3e0,stroke:#ef6c00
    style OUT fill:#e8eaf6,stroke:#3f51b5
    style DETAIL fill:#f3e5f5,stroke:#9c27b0

Реализация MambaBlock¶

class MambaBlock(nn.Module):
    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state

        self.in_proj = nn.Linear(d_model, 2 * d_model, bias=False)
        self.conv = nn.Conv1d(d_model, d_model, d_conv,
                              padding=d_conv-1, groups=d_model)
        self.x_proj = nn.Linear(d_model, d_state * 2 + 1, bias=False)  # B, C, Delta
        self.dt_proj = nn.Linear(1, d_model, bias=True)
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq_len, d_model = x.shape
        x_and_res = self.in_proj(x)
        x, res = x_and_res.split([d_model, d_model], dim=-1)

        x = x.transpose(1, 2)
        x = self.conv(x)[:, :, :seq_len]
        x = x.transpose(1, 2)

        y = self.selective_scan(x)
        return self.out_proj(y * F.silu(res))

Selective Scan (core algorithm)¶

def selective_scan(self, x: torch.Tensor) -> torch.Tensor:
    batch, seq_len, d_model = x.shape

    x_proj = self.x_proj(x)  # [batch, seq_len, d_state*2 + 1]
    B, C, dt = x_proj.split([self.d_state, self.d_state, 1], dim=-1)

    dt = F.softplus(self.dt_proj(dt))  # [batch, seq_len, d_model]

    dA = torch.exp(dt.unsqueeze(-1) * self.A)  # discretized A
    dB = dt.unsqueeze(-1) * B.unsqueeze(-2)    # discretized B

    h = torch.zeros(batch, d_model, self.d_state, device=x.device)
    outputs = []

    for t in range(seq_len):
        h = dA[:, t] * h + dB[:, t] * x[:, t].unsqueeze(-1)
        y = (C[:, t].unsqueeze(-2) * h).sum(-1)
        outputs.append(y)

    return torch.stack(outputs, dim=1)

Mamba-2: State Space Duality (SSD)¶

Paper: "Transformers are SSMs" (Dao & Gu, 2024)

Ключевой инсайт¶

SSM можно представить как умножение на semi-separable matrix: $$Y = \text{SSM}(A, B, C, X) = \text{SS}(A) \cdot (B \otimes X) \cdot C$$

Связь с Attention: - Attention: $Y = \text{softmax}(QK^T/\sqrt{d}) \cdot V$ - SSM: $Y = (\text{SS matrix}) \cdot X$

Оба -- матричные умножения, но SSM структурирован (semi-separable) -> быстрые алгоритмы.

Упрощения¶

Feature	Mamba-1	Mamba-2
Алгоритм	Parallel associative scan	Matrix multiplication
Матрица A	Full matrix	Scalar (shared)
Скорость	1x	2-8x
Качество	Baseline	Same or better

class Mamba2Block:
    def __init__(self, d_model, d_state=128):
        self.A_log = Parameter(1)  # Scalar, not matrix
        self.D = Parameter(d_model)
        self.in_proj = Linear(d_model, d_state * 2 + 1)

    def forward(self, x):
        A = -exp(self.A_log)  # Ensure stability
        BCdt = self.in_proj(x)
        B, C, dt = split(BCdt)
        # SSD computation (hardware-efficient)
        ...

Mamba-3 (ICLR 2026 Oral)¶

Paper: "Mamba-3: Improved Sequence Modeling using State Space Principles" Authors: Lahoti, Li, Chen, Wang, Bick, Kolter, Tri Dao, Albert Gu (Jan 2026)

Три улучшения (Inference-First Design)¶

More Expressive Recurrence -- improved discretization
Complex State Update Rule -- richer state tracking, long-range dependencies
Multi-Input Multi-Output (MIMO) -- лучше использует hardware parallelism при декодинге

Результаты¶

Pareto-frontier для performance under fixed inference budget
Превосходит baselines в retrieval, state-tracking, language modeling
Фокус на inference efficiency (test-time compute scaling)

Model	Инновация	Compute	Memory	Best For
Mamba-1	Selective SSM	O(n)	Constant	Long sequences
Mamba-2	SSD (semi-separable)	O(n)	Constant	Training speed (2-8x)
Mamba-3	MIMO + complex dynamics	O(n)	Constant	Inference efficiency

RWKV¶

RWKV (Receptance Weighted Key Value) -- RNN с Transformer-level качеством и линейной сложностью.

Версии¶

Версия	Год	Инновация
RWKV-4	2023 (EMNLP)	Первая официальная версия
RWKV-5 (Eagle)	2024	Matrix-valued states
RWKV-6 (Finch)	2024	Dynamic recurrence
RWKV-7 (Goose)	Mar 2025	Dynamic State Evolution, surpasses TC0

WKV Formula¶

\[\text{WKV}_t = \frac{\sum_{i=1}^{t-1} e^{w \cdot (i-t)} \cdot K_i \cdot V_i}{\sum_{i=1}^{t-1} e^{w \cdot (i-t)} \cdot K_i}\]

Свойства: - RNN-style: constant memory, sequential generation - Transformer-style: parallelizable training - Без механизма attention

RWKV vs Transformer¶

Aspect	RWKV	Transformer
Compute	10-100x lower (long context)	O(T^2)
Memory	Constant O(1)	Linear O(T) с KV cache
Multilingual	100+ языков (World dataset)	Typically English-centric
Lookback tasks	Weaker	Excellent
Prompt sensitivity	High format sensitivity	Lower

from rwkv.model import RWKV
from rwkv.utils import PIPELINE

model = RWKV(model="RWKV-4-World-7B", strategy="cuda fp16")
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
output = pipeline.generate("Hello", token_count=100, temperature=1.0, top_p=0.7)

Griffin (Google DeepMind, 2024)¶

Paper: "Griffin: Mixing Gated Linear Recurrences with Local Attention"

Gated Linear Recurrence (RG-LRU)¶

\[h_t = a_t \odot h_{t-1} + b_t \odot x_t\]

где $a_t, b_t$ -- input-dependent gates.

Компоненты: 1. RG-LRU: simplified recurrent unit с gating для selective memory 2. Local Attention: sliding window over recent tokens 3. Hybrid: alternation between recurrence и local attention

Гибридные архитектуры (2025-2026)¶

Зачем гибриды?¶

Pure SSMs слабы в: - Copy/recall tasks - Dense global token interactions - Complex multi-hop reasoning

Формула тренда: $$\text{Best Model} = \alpha \cdot \text{Attention} + \beta \cdot \text{SSM} + \gamma \cdot \text{MoE}$$

Jamba (AI21 Labs)¶

Transformer + Mamba + MoE, interleaved layers:

Jamba Block: 7 Mamba + 1 Attention + MoE MLP
Total: 52 layers = 46 Mamba + 6 Attention

Модель	Total params	Active params	Context	Примечание
Jamba (original)	52B	12B	256K	First large-scale hybrid
Jamba 1.5 Large	398B	94B	256K	Production-scale
Jamba Reasoning 3B	--	3B	--	26 Mamba + 2 attention

3x throughput vs comparable pure Transformers.

Другие гибриды¶

Модель	Организация	Особенности
Nemotron-H	NVIDIA	Production-grade, оптимизирован под NVIDIA hardware
Bamba	IBM	Adjustable attention/SSM ratio, enterprise
Hunyuan TurboS	Tencent	Low latency inference
Phi-4-mini-flash	Microsoft	Reasoning + efficient processing

Design Guidelines¶

Компонент	Mamba	Transformer
Local patterns	Best	--
Long-range dependencies	Best	Expensive
In-context learning	Weak	Best
Copy/retrieval tasks	Weak	Best
General reasoning	Good	Good

class HybridBlock(nn.Module):
    """Pattern: 1 attention per N SSM layers (Jamba: 1:7)"""

    def __init__(self, d_model, ssm_ratio=7):
        super().__init__()
        self.ssm_layers = nn.ModuleList([
            MambaBlock(d_model) for _ in range(ssm_ratio)
        ])
        self.attn = AttentionBlock(d_model)
        self.mlp = MLPBlock(d_model)

    def forward(self, x):
        for ssm in self.ssm_layers:
            x = ssm(x) + x  # Residual
        x = self.attn(x) + x
        x = self.mlp(x) + x
        return x

Детали и сравнения¶

Benchmarks¶

Language Modeling (Pile, 3B params):

Model	Perplexity	Inference Speed
Transformer-3B	6.2	1x
Mamba-3B	6.1	3x
Mamba-2-3B	--	5x
RWKV-7-3B	6.3	--

Результат: SSMs матчат Transformers при равном числе параметров.

Extended benchmarks:

Model	Params	MMLU	LongBench	Inference
LLaMA-7B	7B	35.2	28.5	1x
Mamba-3B	3B	32.1	31.2	3x
Mamba-2-3B	3B	33.5	32.8	5x
Jamba-1.5	12B active	52.1	42.5	2x

Memory comparison¶

Context Length	Transformer KV Cache	Mamba State
4K	2 GB	16 MB
32K	16 GB	16 MB
128K	64 GB	16 MB
1M	512 GB	16 MB

Inference Speed¶

Context	Transformer	Mamba	Speedup
2K	Baseline	~1x	--
8K	4x slower	~1x	4x
32K	16x slower	~1x	16x
128K	OOM often	~1x	Massive

Mamba maintains constant inference time regardless of context length.

Training Efficiency¶

Model	Tokens/sec	Memory	Training Time (1T tokens)
LLaMA-7B	1M	80 GB	12 days
Mamba-3B	1.2M	40 GB	10 days
Mamba-2-3B	2M	35 GB	6 days

Long-Context Tasks¶

Task	Transformer	Mamba	Преимущество
Document Q&A (16K)	Good	Excellent	Memory
Genomics (1M tokens)	OOM	Excellent	Length
Audio (100K)	Slow	Fast	Speed
Copy/recall	Excellent	Weaker	Accuracy

Decision Framework¶

Short sequences (<4K)?
  -> Transformer (well-optimized, better reasoning)

Long sequences (>16K)?
  -> Mamba / SSM (linear scaling, constant memory)

Memory-constrained?
  -> Mamba (16 MB state vs GB of KV cache)

Need copy/recall, multi-hop reasoning?
  -> Transformer or Hybrid

General-purpose production LLM?
  -> Hybrid (Jamba-style: 7 Mamba + 1 Attention)

Streaming / real-time?
  -> Mamba / RWKV (O(1) per token)

Формулы (сводка)¶

Формула	Описание
$h_t = \bar{A}h_{t-1} + \bar{B}x_t$	SSM state update
$B_t, C_t, \Delta_t = \text{proj}(x_t)$	Mamba selective parameters
$\bar{A} = \exp(\Delta \cdot A)$	Discretization
$\text{WKV}_t = \frac{\sum \alpha^{t-i} K_i V_i}{\sum \alpha^{t-i} K_i}$	RWKV WKV
$h_t = a_t \odot h_{t-1} + b_t \odot x_t$	Griffin RG-LRU
$Y = \text{SS}(A) \cdot (B \otimes X) \cdot C$	Mamba-2 SSD

Interview Questions¶

1. Что такое SSM и чем отличается от attention?¶

Red flag: "SSM -- это RNN нового поколения"

Strong answer: "SSM сжимает всю историю в fixed-size state через рекуррентное обновление: $h_t = Ah_{t-1} + Bx_t$. Attention вычисляет попарные взаимодействия всех токенов: $O(n^2)$ time, $O(n)$ memory (KV cache). SSM: $O(n)$ time, $O(1)$ memory при инференсе. Ключевое отличие: SSM -- sequential compression (фиксированный state, информация теряется), Attention -- parallel all-to-all (точный доступ к каждому токену). SSM лучше для длинных контекстов, Attention -- для точного retrieval."

2. Что нового в Mamba по сравнению с S4?¶

Red flag: "Mamba -- это просто быстрый SSM"

Strong answer: "S4 использовал фиксированные параметры A, B, C -- модель обрабатывала все токены одинаково. Mamba добавляет selection mechanism: $B_t, C_t, \Delta_t$ зависят от входа. Модель решает что запомнить (ключевые слова), а что пропустить ('the', 'and'). Это dynamic state compression. Плюс hardware-aware parallel scan для эффективного обучения на GPU."

3. Mamba-1 vs Mamba-2 vs Mamba-3?¶

Red flag: "Каждая версия просто быстрее"

Strong answer: "Mamba-1: selective SSM, parallel associative scan, матрица A полная. Mamba-2 (SSD): ключевой инсайт -- SSM = умножение на semi-separable matrix. A стал скаляром, 2-8x ускорение, без потери качества. Mamba-3 (ICLR 2026 Oral): inference-first design. MIMO formulation для лучшего hardware parallelism, complex state dynamics. Pareto-оптимально: лучшее качество при фиксированном inference budget."

4. Почему pure SSM слабее в copy/recall?¶

Red flag: "SSM ещё не доработали"

Strong answer: "Это фундаментальное ограничение архитектуры. SSM сжимает историю в fixed-size state -- информация необратимо теряется. Для copy/recall нужен точный доступ к конкретному прошлому токену, что Attention делает напрямую через $QK^T$. SSM аппроксимирует через сжатый state -- неточно. Поэтому гибриды (Jamba: 7 Mamba + 1 Attention) -- production стандарт: SSM для эффективности, Attention для in-context learning."

5. Спроектируйте архитектуру для 1M-token документов¶

Red flag: "Используем Transformer с Flash Attention"

Strong answer: "Pure Transformer: KV cache = 512 GB при 1M tokens -- нереалистично. Решение -- hybrid: Jamba-style 7:1 ratio (SSM:Attention). Mamba layers для long-range compression ($O(n)$, state = 16 MB). Attention слои через каждые 7 SSM для in-context learning и точного retrieval. MoE MLP для parameter efficiency. Budget-constrained: pure Mamba. Quality-first: Jamba hybrid. Результат: 3x throughput vs pure Transformer, 256K+ контекст."

Ключевые числа¶

Факт	Значение
Transformer complexity	O(n^2 * d)
Mamba complexity	O(n * d^2)
Mamba inference memory	O(d) constant
Mamba-2 speedup vs Mamba-1	2-8x
Mamba state @ any context	16 MB
Transformer KV @ 1M	512 GB
Jamba original	52B total / 12B active
Jamba 1.5	398B total / 94B active
Jamba context	256K tokens
Jamba throughput vs Transformer	3x
RWKV compute advantage (long ctx)	10-100x
Mamba-3	ICLR 2026 Oral, Pareto-frontier

SSM != замена Transformer для всех задач

Mamba-3B матчит Transformer-3B по perplexity, но pure SSM проигрывает на in-context learning, copy/recall, multi-hop reasoning. Тренд 2026: не "SSM vs Transformer", а гибриды (Jamba, Nemotron-H). Pure SSM оправдан для длинных последовательностей (genomics, audio) и memory-constrained сценариев. Для general-purpose LLM -- гибрид.

Заблуждение: O(n) SSM всегда быстрее O(n^2) Transformer

На коротких последовательностях (<4K токенов) Transformer с FlashAttention быстрее Mamba. Причина: constant factor в SSM выше ($O(n \cdot d^2)$ vs $O(n^2 \cdot d)$), а GPU оптимизированы под матричные умножения attention. Mamba выигрывает при длинных контекстах (>16K), где квадратичность начинает доминировать. При 2K токенах speedup Mamba vs Transformer ~1x (паритет).

Заблуждение: SSM state хранит всю историю без потерь

Fixed-size state (16 MB для Mamba-3B) -- это lossy compression. Информация необратимо теряется при сжатии. Для copy task ("повтори 500-й токен") Transformer с KV cache дает точный ответ, а SSM аппроксимирует через сжатый state -- неточно. Именно поэтому гибриды (Jamba) вставляют Attention слои через каждые 7 SSM: для точного retrieval нужен прямой доступ к прошлым токенам.

Самопроверка

SSM state для Mamba-3B = 16 MB при любой длине контекста. Transformer KV cache для LLaMA-7B при 128K = 64 GB. Во сколько раз SSM эффективнее по памяти? Что это означает для batch size при inference?
Jamba использует ratio 7:1 (Mamba:Attention). Почему не 1:1 или 15:1? Что определяет оптимальный ratio?
Реализуйте упрощённый SSM state update ($h_t = Ah_{t-1} + Bx_t$, $y_t = Ch_t$) в PyTorch для 1D входа. Покажите, что при фиксированных A, B, C модель не может быть селективной -- она обрабатывает "the" и ключевые слова одинаково.

Источники¶

Foundational Papers¶

S4 -- Gu et al. (2021) -- arXiv:2111.00396
Mamba -- Gu & Dao (2023) -- arXiv:2312.00752
Mamba-2 -- "Transformers are SSMs" -- Dao & Gu (2024) -- arXiv:2405.21075
Mamba-3 -- "Improved Sequence Modeling" -- ICLR 2026 Oral -- OpenReview:HwCvaJOiCj
RWKV-4 -- Peng et al. (2023) -- arXiv:2305.13048 (EMNLP)
RWKV-⅚ -- Peng et al. (2024) -- arXiv:2404.05892
RWKV-7 -- Peng et al. (2025) -- arXiv:2503.14456

Hybrid Models¶

Jamba -- AI21 Labs (2024) -- arXiv:2403.19887
Griffin -- Google DeepMind (2024) -- arXiv:2402.19427
Nemotron-H -- NVIDIA (2025) -- arXiv:2408.15237
Bamba -- IBM (2025) -- arXiv:2410.05355

Surveys¶

"From S4 to Mamba: Comprehensive Survey on SSMs" -- arXiv:2503.18970 (Mar 2025)
"Mamba-360: Survey of State Space Models" -- ScienceDirect (2025)

Blogs¶

Galileo AI -- "How Mamba Beats Transformers at Long Sequences" (Sep 2025)
AI21 Labs -- "Attention was never enough: Rise of Hybrid LLMs" (Aug 2025)
Maarten Grootendorst -- "Visual Guide to Mamba & SSMs"
Princeton PLI -- "Mamba-2: Algorithms and Systems"

Формула	Описание
\(h_t = \bar{A}h_{t-1} + \bar{B}x_t\)	SSM state update
\(B_t, C_t, \Delta_t = \text{proj}(x_t)\)	Mamba selective parameters
\(\bar{A} = \exp(\Delta \cdot A)\)	Discretization
\(\text{WKV}_t = \frac{\sum \alpha^{t-i} K_i V_i}{\sum \alpha^{t-i} K_i}\)	RWKV WKV
\(h_t = a_t \odot h_{t-1} + b_t \odot x_t\)	Griffin RG-LRU
\(Y = \text{SS}(A) \cdot (B \otimes X) \cdot C\)	Mamba-2 SSD