Проблема квадратичной сложности трансформеров
Классические языковые модели на базе архитектуры Transformer используют механизм self-attention, который требует вычисления связей между каждым токеном и всеми остальными токенами в контексте. Это приводит к квадратичной вычислительной сложности O(N²), где N — длина последовательности. Кроме того, модель хранит KV-кэш (ключи и значения) для каждого токена, что вызывает линейный рост потребления видеопамяти. Для длинных текстов (книги, код, длинные диалоги) это становится узким местом, ограничивающим масштабируемость.
Решение Mamba: Селективные пространства состояний (S4)
Исследователи Альберт Гу и Три Дао предложили архитектуру Mamba, основанную на State Space Models (SSM), но с ключевым улучшением — селективностью. В отличие от классических RNN, которые забывали информацию на длинных дистанциях, и фиксированных SSM, которые не могли адаптироваться к контенту, Mamba динамически обновляет параметры для каждого токена.
Модель использует «механизм отбора» (selection mechanism), где параметры A, B и C вычисляются на лету в зависимости от входного токена. Это позволяет модели:
- Запоминать важную информацию надолго (высокий вес B).
- Сбрасывать шум или нерелевантные данные (низкий вес B).
Это достигается без сравнения всех токенов со всеми, сохраняя линейную сложность обработки.
Параллелизм и скорость обучения
Главный недостаток рекуррентных архитектур — невозможность параллельного обучения из-за зависимости текущего шага от предыдущего. Mamba решает эту проблему, используя свойство ассоциативности операций обновления состояния. Это позволяет применять алгоритм scan (или parallel scan), который разбивает последовательность на части, вычисляет их параллельно на GPU, а затем объединяет результаты. В результате Mamba обучается так же быстро, как трансформеры, но требует меньше памяти при инференсе.
Сравнительные характеристики
Ниже приведено сравнение ключевых метрик архитектуры Mamba и классического Transformer:
| Характеристика | Transformer (Self-Attention) | Mamba (S4 + Selection) |
|---|---|---|
| Вычислительная сложность | O(N²) (квадратичная) | O(N) (линейная) |
| Потребление памяти (KV-кэш) | Растет линейно с длиной контекста | Постоянная (фиксированное состояние) |
| Скорость инференса на длинных контекстах | Замедляется пропорционально N | Почему не зависит от длины (постоянная скорость) |
| Механизм внимания | Глобальное внимание (все ко всем) | Селективное внимание (вход-зависимые параметры) |
| Параллелизм при обучении | Полный | Полный (благодаря parallel scan) |
Почему это важно
Появление Mamba демонстрирует, что архитектура, не использующая механизм внимания, может конкурировать с трансформерами по качеству генерации текста, но превосходить их по эффективности. Это открывает путь к созданию моделей, способных работать с контекстом в сотни тысяч токенов на потребительском оборудовании, что критически важно для обработки больших документов, кодовых баз и долгих историй.
Источник: Habr ↗