tgoop.com/quant_prune_distill/172
Last Update:
Метод
State-space модель по существу - это рекуррентная модель вида (в дискретной форме):h_{t} = A h_{t-1} + B x_{t}
y_{t} = C h_{t}
Где x - вход, h - скрытое состояние, y - выход. A, B, C - обучаемые матрицы параметров, параметризованные специальным образом. Есть еще шаг дискретизации и преобразование из непрерывной формулировки в дискретную с некоторым шагом дискретизации \Delta. В отличие от типичной рекуррентной архитектуры а-ля vanilla RNN, LSTM между h и y нет нелинейности. Благодаря этому формулу можно развернуть и представить в виде свертки, которая эффективно параллезируется на GPU.
Важно, что проекции A, B, C не зависят от времени, и динамика состояний сети обладает свойством Linear Time Invariance (LTI). Данное свойство накладывает ограничения на возможность моделирование определённых зависимостей и авторы рассматривают пару синтетических задач, которые принципиально не решаются state-space моделями с LTI:
1⃣️️️️️️ Selective Copying, где нужно запомнить и вывести несколько элементов последовательности с произвольным расстоянием между запоминаемыми элементами.
2⃣️️️️️️ Induction Heads, где приняв начало словосочетания на вход, на основе полученной ранее информации - продолжить. Если на входе Harry, выдать следует Potter (прочитай другую книгу…).
Данные задачи требуют от модели способности работать с неоднородными временными последовательностями и адаптивно учитывать контекст, которой нет в state-space моделях со статичными сверточными ядрами.
Авторы вводят два понятия с похожим звучанием, но разным значением:
🔵️️️️️️ Efficiency (компактное внутреннее состояние и дешевизна вычислений)
🔵️️️️️️ Effectiveness (способность хранить релевантную информацию)
Трансформер effective - так как может хранить в себе (во всяком случае в теории) контекст произвольной длины. Но не efficient, ибо состояние (key и valuе кэш) растет пропорционально длине последовательности, что делает инференс трансформера дорогим по памяти и вычислениям. Причем большинство хранимой информации избыточно.
Рекуррентные сети efficient, ибо хранят состояние постоянного размера, не зависящего от длины последовательности. Но в это маленькое состояние не всегда удается на практике впихнуть всю релевантную информацию.
Идеальная модель, должна потому адаптивно отбирать нужный для задачи контекст.
Чтобы добиться данного свойства, авторы предлагают делать параметры state-space модели зависящими от входа - проекции A, B, C, как и шаг \Delta. Последнее позволяет заодно учитывать и временную неоднородность.
Развернуть рекуррентную зависимость как для LTI state-space модели теперь не получится. Вместо этого авторы реализуют эффективный алгоритм parallel scan на кастомных чудо-kernel дающий высокую производительность.
Чтобы иметь возможность запихнуть много информации в скрытое состояние, используют большой expansion фактор внутри SSM (10-100). Современные GPU зачастую memory-bound, потому чтобы не тратить драгоценные миллисекунды на перекачку данных из HBM в кэш GPU, это расширенное состояние материализуется только на GPU.
И в духе Flash Attention (один из авторов - автор Flash Attention) чтобы дополнительно сэкономить на памяти и ее трансфере, промежуточные состояния отправляются в /dev/null после прямого прохода и пересчитываются снова на обратном проходе.
Далее авторы, предлагают полностью однородную архитектуру, состоящую только из Mamba блоков, в то время как трансформер - перемежающиеся attention блоки и MLP, и прошлые state-space модели состояли собственно из state-space слоев и gated MLP.
Так как алгоритм впридачу к свойствам S4 новый тип state-space слоев обладает селективностью (selective) и использует сканирование (scan) для краткости авторы обозначают новый тип архитектуры - S6.
BY КПД
Share with your friend now:
tgoop.com/quant_prune_distill/172