State Space Models (SSM)

Модель пространства состояний (SSM) преобразует входную последовательность в выходную с помощью вектора скрытого состояния, обновляемого на каждом шаге, поэтому объём памяти не меняется с увеличением длины последовательности. Эта идея пришла из теории управления, где линейную систему записывают как обновление состояния h' = Ah + Bx и выход y = Ch + Dx и где фильтр Калмана с 1960 года оценивает такие состояния по зашумлённым измерениям. SSM в глубоком обучении учатся определять A, B и C по данным и объединяют множество таких слоёв. Начиная с S4 в 2021 году и Mamba в 2023 году, они конкурируют с трансформерами при обработке длинных последовательностей текста, аудио и данных временных рядов, причём вычислительные затраты растут линейно с длиной последовательности.

Как работают SSM#

Слой нейронной SSM преобразует непрерывные уравнения в дискретную рекуррентную формулу с обучаемым шагом Δ, используя правило удержания нулевого порядка Ā = exp(ΔA). На каждом шаге вычисляются h_t = Ā h_(t-1) + B̄ x_t и y_t = C h_t. Если A, B и C не меняются со временем, разворачивание рекуррентной формулы показывает, что выход представляет собой свёртку входа с ядром (CB̄, CĀB̄, CĀ²B̄, ...). Поэтому при обучении тот же слой может работать как свёртка, параллельно обрабатывая всю последовательность на GPU, а при генерации — как рекуррентная модель, выполняя одно обновление состояния постоянного размера на токен. В PyTorch оба варианта дают одинаковые выходы:

import torch

torch.manual_seed(0)
L, N = 64, 16  # длина последовательности, размер состояния
x = torch.randn(L)

# Непрерывная диагональная SSM, дискретизированная с размером шага delta (удержание нулевого порядка)
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# Рекуррентный вариант: одно обновление состояния на шаг, как при генерации
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# Свёрточный вариант: предварительное вычисление ядра K_k = C . (A_bar^k * B_bar), как при обучении
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Код выводит Same outputs: True. На каждом шаге рекуррентная формула хранит только 16 чисел состояния, а предварительно вычисленное ядро позволяет при обучении параллельно вычислить все 64 выхода.

От S4 к Mamba#

Первые глубокие SSM было слишком медленно и затратно по памяти обучать в масштабе. S4, созданная Альбертом Гу, Караном Гоэлом и Кристофером Ре, наделила матрицу состояния структурой, удешевляющей вычисление ядра свёртки, и инициализацией, рассчитанной на сохранение длинной истории. Она установила наилучшие результаты на всех задачах теста Long Range Arena, стала первой моделью, решившей задачу Path-X с последовательностями длиной 16K шагов, и генерировала данные в 60 раз быстрее трансформера.

Фиксированное ядро одинаково обрабатывает все входные данные, поэтому не может пропустить лишнюю информацию и сохранить важный факт. Mamba, созданная Альбертом Гу и Три Дао, делает Δ, B и C функциями текущего токена, позволяя модели выбирать, что записывать в состояние, а что забывать. Эти селективные параметры не позволяют использовать свёрточное представление, поэтому Mamba вычисляет рекуррентную формулу с помощью параллельного сканирования, учитывающего особенности аппаратуры. Пропускная способность Mamba при инференсе в пять раз превысила показатели трансформеров; Mamba-3B превзошла трансформеры того же размера и сравнялась с моделями вдвое крупнее. Каждый блок Mamba дополняет селективную SSM короткой свёрткой и вентилем SiLU:

Mamba-2 показала, что селективные SSM и одна из форм линейного внимания — два представления одних и тех же вычислений. Авторы называют эту концепцию дуальностью пространства состояний; она позволила ускорить основной слой в 2–8 раз.

SSM и смежные архитектуры#

• Трансформеры: Self-attention сопоставляет каждый токен со всеми остальными, поэтому может точно скопировать любой предыдущий токен, но вычислительные затраты растут квадратично, а KV-кэш растёт вместе с контекстным окном. SSM сжимает историю в состояние фиксированного размера — это дешевле, но приводит к потере информации. В статье 2024 года Repeat After Me показано, что трансформеры лучше SSM с фиксированным состоянием копируют и извлекают данные из длинных входных последовательностей.
• Рекуррентные нейронные сети (RNNs): RNN или LSTM выполняет нелинейное обновление, поэтому при обучении приходится последовательно обрабатывать каждый токен. Обновления SSM линейны по состоянию, что позволяет использовать свёртку или параллельное сканирование.
• Скрытые марковские модели (HMMs): HMM имеет конечный набор дискретных состояний с вероятностями переходов. Линейная гауссовская модель пространства состояний, лежащая в основе фильтра Калмана, — её непрерывный аналог, и глубокие SSM сохраняют этот непрерывный вектор состояния.

Где применяются SSM#

Большинство промышленных решений используют гибридные большие языковые модели (LLMs), в которых оставлено несколько слоёв внимания для точного извлечения данных. Jamba от AI21 чередует слои трансформера, Mamba и смеси экспертов (MoE), помещается на один GPU с 80 ГБ памяти и показывает хорошие результаты при контекстах длиной до 256K токенов. Гибридные модели Granite 4.0-H от IBM сочетают Mamba-2 и слои трансформера в соотношении 9:1. Vision Mamba применяет селективное сканирование к патчам изображения, а SSM также моделируют необработанное аудио и последовательности ДНК, слишком длинные для полного внимания. Классический вариант используется в трекере Ultralytics YOLO26: трекер TrackTrack по умолчанию, как ByteTrack и BoT-SORT, предсказывает следующую рамку каждого объекта с помощью фильтра Калмана. Его восьмимерное состояние содержит центр и размер рамки, а также их скорости; состояние обновляется матрицей перехода с постоянной скоростью и линейной матрицей наблюдения, которая считывает четыре координаты рамки.