Обо мне

Привет, меня зовут Василий Техин. За последние 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. Сжатое состояние: вся история последовательности сжимается в вектор фиксированного размера.

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

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

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:

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. Как селективность помогает анализировать ДНК?