Efficiently Modeling Long Sequences with Structured State Spaces [2021] - как дотер стал нейросетью

Для тех, кто не знает - я не особый любитель длинных математических статей. В целом, я не умею с адекватной скоростью читать и воспринимать много линала. Наверняка в телеграме существует большое количество умных постов про S4 с кратким пересказом его математики, и если вы из тех, кто способен такое воспринимать, поздравляю - данный пост не для вас.

Я постарался, вооружившись гайдом, уловить основной смысл данной архитектуры, где она находится по отношению с известными широким кругам. Итак, поехали.

Представим, что существует "ячейка памяти" - хранилище-вектор, который обновляется с учётом предыдущего состояния ячейки, последнего входа и каких-то обучаемых параметров. Помимо памяти есть функция выхода, которая берёт новое состояние памяти, последний вход и выдаёт выход наружу.

Мы уже знаем реализации подобных абстракций. Простейшая RNN, GRU/LSTM - все мы их любим, но у них есть жирная проблема - их нужно считать шаг за шагом, а значит, нельзя применить много компьюта и обработать кучу информации за раз, так, как это умеют трансформеры, но сами трансформеры фэйлятся на огромных контекстах.

Итак, помимо RNN и GRU существует State Space Model - ещё один формат ячейки памяти (в его основе всего лишь парочка матричных умножений), но у него есть крутая особенность. Вычисление рода "прогнать SSM на последовательности", оказывается, можно переформулировать в другую функцию - свёртку, для которой можно предпосчитать веса. При добавлении ещё одного фокуса (FFT) эту свёртку можно считать быстрее, чем втупую, что в итоге позволяет по сути быстро применять SSM на всей последовательности.

Далее, у SSM есть 2 проблемы - они херово работают, и хвалёное "быстро посчитать" на самом деле не такое уж и быстрое.

Чтобы решить первое, был придуман магический гиппопотам - инициализация одной из матриц внутри SSM таким образом, чтобы она была изначально ближе к пространству чего-то разумного.

Вкратце, вторая проблема заключается в том, что для подсчёта весов свёртки нужно умножать много матриц, а нам вообще-то лень - бумага нынче дорогая. Для этого придумывают магический костыль - Diagonal Plus Low-Rank. Я не стал разбираться в деталях, если вам интересно, отсылаю к разбору, но одну из матриц просто (нихера не просто в реальности) представляют не как обучаемую матрицу весов, а как результат операций над другими обучаемыми сущностями.

В результате, объединив описанные хаки, и получается S4 - хитрая и быстрая вариация "RNN"-ки, которую успешно применяют на сверхдлинных последовательностях.

Замечу, что это не первая статья, которую я обозреваю, в которой засчёт убирания нелинейностей удаётся всё очень сильно ускорить и упростить - напомню про RetNet. Нет никаких гарантий, что "мощность" архитектуры достигается как раз засчёт этих нелинейностей.

Кроме того, скажу честно - я банально не верю, что прорывные архитектуры будут основаны на какой-то сложной математике. Через пару лет окажется, что есть какая-нибудь суперпростая штука, которая делает всё то же самое даже лучше. Это не исключает, что математика будет вдохновлять на прогресс и в какой-то момент натолкнёт ресёрчеров на нечто крутое, но само это крутое будет очень простым.

@knowledge_accumulator