Часть 4: Mamba — State Space Models vs трансформеры

Обо мне

Привет, меня зовут Василий Техин. За последние 6 лет я прошёл путь от новичка в ML который ни в чем не разбирается, до человека, который может помочь разобраться другим не используя сложной математике и приправляя обьяcнение наглядными примерами. В этой серии статей я разбираю прорывные модели "на пальцах":

1. ResNet-18 — Архитектура, покорившая глубину

2. Vision Transformer (ViT) — Когда трансформеры научились видеть

3. Diffusion Transformer (DiT) — Stable Diffusion 3 как она есть Сегодня: Mamba — революция в обработке длинных последовательностей!

---

Пролог: Почему Mamba?

Представьте тусовку ученых исследователей, где каждый должен переговорить со всеми. Для 10 человек — 100 диалогов (да, мы считаем, что разговоры первый со вторым и второй с первым отличаются и да да мы будем считать, что он и сам с собой говорит, кто поймет этих ученых!). Для 1000 — 1 000 000! Это квадратичная сложность (O(n²)) внимания в трансформерах.
Mamba (State Space Models) решает это:

• Замена внимания на линейную рекуррентную сеть (O(n)).

• Селективность: динамический выбор важной информации.

• Ускорение в 5 раз на длинных текстах.

Сравнение временных сложностей

---

Ключевые компоненты Mamba

Архитектура мамбы

1. State Space Models (SSM)

State Space Models (SSM) — семейство линейных моделей для последовательностей, основанных на дифференциальных уравнениях. Их ключевые особенности:

1. Линейность: состояние обновляется линейно (s(t) = A·s(t-1) + B·x(t)), без нелинейных активаций.

2. Дискретизация: параметры выводятся из непрерывных уравнений для работы с дискретными токенами.

3. Сжатое состояние: вся история последовательности сжимается в вектор фиксированного размера.

Отличия от других моделей

Модель	Ключевые особенности	Отличие от SSM
RNN/GRU/LSTM	Нелинейные активации (tanh, sigmoid)	SSM — полностью линейны
Трансформеры	Внимание (O(n²))	SSM — рекуррентные (O(n))
S4 SSM	Линейность + инвариантность во времени	Mamba добавляет селективность (B, C, Δt зависят от входа)
Mamba	SSM + селективность + параллельный scan	Динамическая адаптация к важности токенов

Ядро архитектуры:

ds(t)/dt = A·s(t) + B·x(t)
y(t) = C·s(t) + D·x(t)

Где:

• x(t) — входной сигнал (токен) в момент времени t

• s(t) — скрытое состояние (память модели)

• ds(t)/dt — скорость изменения состояния

• A — матрица переходов (определяет "забывание" прошлого)

• B — матрица входа (влияние нового токена на состояние)

• C — матрица выхода (преобразует состояние в выход)

• D — сквозная связь (прямое влияние входа на выход)

• y(t) — выходной сигнал

Дискретизация для токенов:

A_d = exp(-Δt * A)   # Дискретная матрица переходов
B_d = (A⁻¹)(exp(Δt·A) - I)B  # Дискретная матрица входа
s_t = A_d · s_{t-1} + B_d · x_t

Где:

• A_d — дискретный аналог матрицы A (определяет "сохранение памяти")

• B_d — дискретный аналог матрицы B (влияние входа на обновление состояния)

• Δt — шаг дискретизации (учится: большой = долгая память, маленький = быстрое забывание)

2. Селективность

Уникальность Mamba:

B = Linear(x_t)  # Линейное преобразование входа
Δt = softplus(Linear(x_t))  # Активация softplus: log(1 + exp(x))

Где:

• Linear() — полносвязный слой (y = Wx + b)

• softplus() — гладкая версия ReLU, гарантирует Δt > 0

Модель динамически решает:

• "Важный токен? Запомню надолго (↑ Δt)"

• "Шум? Проигнорирую (↓ Δt)"

3. Параллельный Scan

Решение проблемы RNN:

state = cumprod(A) * cumsum(B)  # O(log n) вместо O(n)

Где:

• cumprod() — кумулятивное произведение (a1, a1a2, a1a2*a3, ...)

• cumsum() — кумулятивная сумма (b1, b1+b2, b1+b2+b3, ...)

Как работает:

1. Разбиваем последовательность на блоки

2. Вычисляем промежуточные состояния внутри блоков параллельно

3. Объединяем результаты между блоками

4. Повторяем иерархически (дерево слияний)

---

Разберём на простом примере

Обработка фразы "Мама мыла раму"

1. Токенизация

"М"→10, "а"→11, " "→12, "м"→13, "ы"→14, "л"→15, "р"→16, "у"→17
Последовательность: [10, 11, 13, 11, 12, ...]

2. Эмбеддинг

Каждый токен → вектор:

10 → [0.2, -0.1, 0.4, 1.1]
11 → [0.7, 0.3, -0.2, 0.5]

3. Блок Mamba

Для токена "м" (позиция 2):

1. Локальный контекст (Conv1D):

   conv("Ма") = [0.5, 0.1, 0.3, 0.9]  # Смешение соседей
   

2. Селективное обновление состояния:

   Δt = sigmoid("м") → 0.7  # Высокая важность
   new_state = (1 - 0.7*A) * state + 0.7*B * x_token
   

3. Gating механизм:

   gate = sigmoid(Linear(x))  # Решает, какую информацию передать
   output = gate * (C * state + D * x_token)
   

4. Выход: y = output

4. Предсказание

После 12 блоков:

logits = linear_layer(output)  # Вероятности символов
next_token = argmax(logits) → "а"  # "Мама..."

---

Код: Сердце Mamba

Ядро SSM

class MambaBlock(nn.Module):
    def forward(self, x):
        # Локальный контекст
        x_conv = depthwise_conv1d(x)
        
        # Селективные параметры
        Δ, B, C = self.x_proj(x_conv)  
        A_d = exp(-Δ * self.A)
        
        # Параллельный scan
        state = cumprod(A_d) * cumsum(B * x_conv)
        y = einsum('btd,btd->bt', state, C) + self.D * x_conv
        
        # Gating механизм
        return y * torch.sigmoid(self.gate_proj(x))

Параллельный Scan

def parallel_scan(A, B):
    A_cum = torch.cumprod(A, dim=1)  # Продукт: [a0, a0*a1, a0*a1*a2, ...]
    B_cum = torch.cumsum(B * A_cum, dim=1)  # Взвешенная сумма
    return B_cum / (A_cum + 1e-7)  # Нормализация

---

Сравнение: Mamba vs Трансформер

Критерий	Трансформер	Mamba
Сложность	O(n²)	O(n)
Память (инференс)	O(n)	O(1)
Длина контекста	~128K	>1M токенов
Качество	SOTA	Сопоставимо
Селективность	Жёсткое внимание	Динамический выбор

Сильные стороны Mamba:

1. Анализ ДНК/белковых последовательностей.

2. Работа на edge-устройствах (телефоны).

Слабые стороны:

1. Короткие тексты (трансформеры эффективнее).

2. Мультимодальные задачи (пока).

---

Примеры применения

1. 🧬 BioMamba:

   • Анализ генома человека (3 млрд пар оснований).

2. 🖼️ Vim (Vision Mamba):

   • На 30% быстрее ViT на ImageNet.

   • Статья.

Резюме

Mamba — прорыв для длинных последовательностей:

• State Space Model с селективностью.

• Сложность O(n) vs O(n²) у трансформеров.

• Применения: геномика, длинные тексты, edge-устройства.

• Ограничения: мультимодальность, короткие контексты.

Ссылки:

1. Оригинальная статья

2. Объяснение SSM

---

Философский итог

Mamba не заменяет трансформеры, а расширяет инструментарий ИИ:

• Для длинных данных — оптимальный выбор.

• Для мультимодальности — трансформеры сохраняют преимущество. Будущее: Гибридные архитектуры, где Mamba обрабатывает историю, а внимание — кросс-модальные связи.

---

Проверь себя

1. Почему сложность Mamba O(n), а не O(n²)?

2. Как селективность помогает анализировать ДНК?