Как работает Mamba (State Space Model) и чем она лучше трансформера?

Краткий тезис

Mamba — это архитектура нейросети, основанная на State Space Model (SSM) вместо механизма attention. В отличие от трансформера с квадратичной сложностью O(n²), Mamba обрабатывает последовательности за линейное время O(n) благодаря рекуррентному обновлению скрытого состояния. Ключевое нововведение — selective scan, который делает параметры SSM зависимыми от входных данных, что позволяет модели динамически выбирать, какую информацию запоминать. Это даёт возможность работать с контекстами длиной более 1 миллиона токенов без роста памяти, хотя на сложных задачах рассуждения качество Mamba пока немного уступает трансформеру.

1. Что такое State Space Model (SSM)

State Space Model — это математическая модель, описывающая динамическую систему через state|скрытое состояние h(t), которое эволюционирует во времени. В контексте обработки последовательностей (текстов, сигналов) SSM преобразует входной сигнал x(t) в выходной y(t) через систему дифференциальных уравнений.

В дискретном виде (для токенов) базовое уравнение SSM выглядит так:

h_t = A * h_{t-1} + B * x_t
y_t = C * h_t

Где:

• h_t — state|скрытое состояние в момент t (вектор размерности d_state).
• x_t — входной токен (эмбеддинг).
• y_t — выход для токена t.
• A — матрица перехода состояния (d_state × d_state), определяет, как состояние затухает/эволюционирует.
• B — матрица проекции входа на состояние (d_state × d_in).
• C — матрица проекции состояния на выход (d_out × d_state).

В классических SSM (например, S4) матрицы A, B, C — линейные и статические (не зависят от входа). Это ограничивает способность модели адаптироваться к контексту.

2. Отличие SSM от Attention

Attention в трансформере вычисляет попарные веса между всеми токенами последовательности:

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

Сложность — O(n²) по длине последовательности n, так как нужно вычислить матрицу внимания n×n.

SSM же обрабатывает последовательность рекуррентно: каждый шаг зависит только от предыдущего состояния и текущего входа. Сложность — O(n) по времени и O(1) по памяти (хранится только состояние).

3. Selective Scan — ключевая инновация Mamba

Проблема классических SSM (S4, S5): матрицы A, B, C фиксированы, поэтому модель не может «выборочно» игнорировать неважные токены или запоминать важные. Трансформер решает это через attention — веса зависят от содержимого.

Mamba вводит selective scan: параметры B и C (а также шаг дискретизации Δ) становятся функциями от входного токена x_t:

B_t = Linear_B(x_t)
C_t = Linear_C(x_t)
Δ_t = softplus(Linear_Δ(x_t))

Теперь:

• Если токен неважен, B_t может быть близок к нулю → состояние почти не обновляется.
• Если токен важен, B_t активируется → информация записывается в состояние.
• Δ_t управляет скоростью обновления (аналог gate в LSTM).

Это позволяет Mamba динамически выбирать, какую информацию сохранять в скрытом состоянии, подобно тому, как attention выбирает, на какие токены смотреть.

Реализация selective scan в коде (упрощённо):

import torch
import torch.nn.functional as F

def selective_scan(u, delta, A, B, C, D):
    """
    u: входные токены [batch, length, d_in]
    delta: шаг дискретизации [batch, length, d_in]
    A: матрица состояния [d_state, d_state]
    B: проекция входа [batch, length, d_state]
    C: проекция выхода [batch, length, d_state]
    D: skip-connection [d_out]
    """
    batch, length, d_in = u.shape
    d_state = A.shape[0]
    h = torch.zeros(batch, d_state, device=u.device)
    ys = []
    for t in range(length):
        # дискретизация: A_bar = exp(delta * A)
        A_bar = torch.exp(delta[:, t, :].unsqueeze(-1) * A.unsqueeze(0))  # [batch, d_state, d_state]
        B_bar = delta[:, t, :].unsqueeze(-1) * B[:, t, :].unsqueeze(-2)   # [batch, d_state, d_in]
        # обновление состояния
        h = A_bar @ h.unsqueeze(-1) + B_bar @ u[:, t, :].unsqueeze(-1)
        h = h.squeeze(-1)
        # выход
        y = (C[:, t, :] @ h.unsqueeze(-1)).squeeze(-1) + D * u[:, t, :]
        ys.append(y)
    return torch.stack(ys, dim=1)

На практике используется параллельный алгоритм (ассоциативный сканер), чтобы ускорить обучение.

4. Вычислительная сложность и память

Трансформер:

• Обучение: O(n²) на слой из-за attention.
• Инференс: O(n²) при генерации (нужно пересчитывать attention для каждого нового токена с учётом всей истории). Используют KV-cache, но он растёт линейно с длиной контекста.

Mamba:

• Обучение: O(n) на слой (selective scan можно распараллелить через ассоциативный сканер, но всё равно линейно).
• Инференс: O(1) на шаг (только обновление состояния). Память константна — не нужно хранить историю.

Сравнение для контекста 1M токенов:

• Трансформер: attention map заняла бы ~1 ТБ памяти (при d=4096). Даже с FlashAttention — десятки ГБ.
• Mamba: состояние ~ d_state × d_state (например, 16×16 = 256 чисел) — несколько килобайт.

5. Преимущества Mamba перед трансформером

1. Линейная сложность — позволяет обрабатывать сверхдлинные последовательности (1M+ токенов) без переполнения памяти.
2. Константная память на инференсе — идеально для агентов, которые работают с длинной историей диалога или документами.
3. Быстрый инференс — генерация каждого токена требует O(1) операций, а не O(n).
4. Меньше параметров — при равном качестве Mamba часто компактнее трансформера (нет параметров attention).
5. Хорошая производительность на задачах с длинным контекстом — например, суммаризация книг, анализ логов, обработка геномных последовательностей.

6. Недостатки и ограничения

1. Качество на сложных рассуждениях — на бенчмарках вроде GSM8K, MATH, MMLU Mamba уступает трансформеру сопоставимого размера. Причина: рекуррентное состояние — «бутылочное горлышко»; модель не может напрямую сопоставить далёкие токены.
2. Параллелизация обучения — хотя selective scan можно ускорить, он менее эффективен на GPU, чем матричные умножения attention (особенно с FlashAttention).
3. Гибридные подходы — лучшие результаты дают гибриды (Mamba + attention), например Jamba, Zamba.
4. Меньше экосистемы — меньше предобученных моделей, инструментов, оптимизаций, чем у трансформеров.

7. Когда выбирать Mamba, а когда трансформер

8. Mamba в контексте RAG и AI-агентов

В Agentic RAG агент часто работает с длинными контекстами: история диалога, несколько документов, результаты вызовов инструментов. Трансформеры с attention быстро упираются в лимит контекста (даже 128k — дорого). Mamba позволяет:

• Хранить всю историю в скрытом состоянии без потери информации.
• Делать быстрые итеративные запросы к базе знаний.
• Обрабатывать потоки данных (например, логи в реальном времени).

Пример: агент, который анализирует 500-страничный отчёт и отвечает на вопросы. С Mamba можно загрузить весь отчёт как последовательность токенов и задавать вопросы без перезагрузки контекста.

Пет-проект для закрепления

Задача: Реализовать упрощённую версию Mamba для задачи классификации текстов с длинным контекстом (например, определение тональности по целой книге).

Инструменты: Python, PyTorch, библиотека mamba-ssm (официальная реализация).

Шаги:

1. Установить mamba-ssm и causal-conv1d.
2. Загрузить датасет длинных текстов (например, Amazon Reviews с полными текстами).
3. Построить модель: слой Mamba (d_model=256, d_state=16) + голову классификации.
4. Обучить на задаче бинарной классификации (позитивный/негативный отзыв).
5. Сравнить с трансформером (например, DistilBERT) по скорости инференса и качеству на длинных текстах (>512 токенов).

Ожидаемый результат: Mamba покажет сопоставимое качество, но будет в 2-5 раз быстрее на инференсе и сможет обрабатывать тексты любой длины без ограничения.

Связь с другими вопросами

Навигация

• Предыдущий: 712
• Следующий: 714
• Индекс: 00. Индекс разборов