Mamba: Linear-Time Sequence Modeling with Selective State Spaces Albert Gu, Tri Dao Статья: https://arxiv.org/abs/2312.00752 Код: https://github.com/state-spaces/mamba Твиттер-тред: https://twitter.com/_albertgu/status/1731727672286294400 Свежее продолжение истории про state space models (SSM), а точнее structured SSM или S4 (https://t.center/gonzo_ML/1424). S4 имеет рекуррентную формулировку, к тому же её можно реализовать и через свёртку, имея линейную или около того сложность от длины входной последовательности. Модели этого класса продемонстрировали хорошие результаты в моделировании длинных последовательностей, и, конечно, все давно ждут, когда уже наконец мы побьём трансформеры на моделях большого размера на текстовых задачах. Пока это не очень получалось, основные крутые результаты были на непрерывных сигналах типа аудио и зрения. Текущая работа предлагает новый класс selective state space models и закрывает этот пробел, получая качество трансформеров с линейным масштабированием по размеру входа. Напомним, что S4 задаётся четырьмя параметрами: A,B,C и ∆, которые определяют двухэтапную seq2seq трансформацию, где вход x(t) трансформируется в скрытое состояние h(t), а оно в свою очередь в выход y(t). В новой работе наконец пришли к стандартным обозначениям входа и скрытого состояния, а не как в работе про S4, где вход был u(t), а скрытое состояние x(t). Если посмотреть на рекуррентную реализацию, то это выглядит так: h_t = Ah_{t−1} + Bx_t y_t = Ch_t На первом этапе непрерывные параметры ∆, A, B дискретизуются по заданному правилу, а на втором происходит вычисление либо через линейную рекуррентность, либо через глобальную свёртку. Рекуррентность хороша для инференса, свёртка хорошо параллелится и удобна для обучения. Модель обладает свойством Linear Time Invariance (LTI), её динамика постоянна во времени. Благодаря этому свойству модель можно эффективно вычислить свёрткой. Текущая работа демонстрирует фундаментальные ограничения LTI и челлендж эффективной реализации. Благодаря структуре в матрицах параметров, каждая из них (A, B, C) может быть представлена N числами. Для обработки входной последовательности x длины L с D каналами и с размером батча B, SSM применяется к каждому каналу независимо, и общее скрытое состояние имеет размерность DN. Работа по всему входу соответственно требует O(BLDN) памяти и вычислений. По мнению авторов, фундаментальная проблема моделирования последовательностей заключается в компрессии контекста в меньшего размера состояние. На трейдофы популярных моделей можно смотреть с этой точки зрения. Механизм внимания в этом смысле effective (позволяет получать хороший результат), но inefficient (требует слишком много ресурсов). Неэффективность его от того, что не происходит сжатия контекста -- весь контекст в виде KV кеша явно хранится для инференса, отсюда он в трансформерах линейный по времени, отсюда же и квадратичное обучение. Рекуррентные модели наоборот efficient -- у них фиксированного размера состояние, отсюда и инференс за константное время и линейное по времени обучение. Но качество результата сильно зависит от того, насколько хорошо состояние хранит в себе контекст. Показали это на двух модельных задачах, требующих понимания контекста, где недостаточно константной динамики. Одна задача -- это Selective Copying, модификация обычного Copying, где расстояние между запоминаемыми токенами может варьировать и модели надо выборочно запоминать или игнорировать входные данные в зависимости от их содержимого. Другая задача -- Induction Heads из Transformer Circuits. Там надо делать prefix matching внутри контекста и далее копирование. Для LTI систем эти задачи фейлятся. В итоге, авторы считают, что фундаментальный принцип для построения sequence models -- это selectivity, контекстно-зависимая способность фокусироваться или отфильтровывать входы в состояние последовательности. Авторский метод решения этой проблемы -- позволить параметрам, отвечающим за взаимодействие с последовательностью (это ∆, B, C), зависеть от входа (здесь через линейные проекции, но возможны и иные варианты).