Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
[Статья][Код]

Современные нейронные сети, обрабатывающие пространственно-временные данные различной природы будь то текст 📝, изображения 📷, аудио 🎵 и видео 📹 так или иначе обладают механизмом перемешивания каналов (channel mixing), обрабатывающим независимо признаки для каждого элемента последовательности, и механизмом обработки последовательности (sequence mixing), использования взаимосвязей между элементами.

В сегодняшнем рассказе речь пойдет про sequence mixing.

Существуют разнообразные опции sequence mixing. Операция смешивания может не зависеть от входа, как например свертка или обучаемая матрица L x L (L - длина последовательности) в MLP-Mixer, S4 и H3 state-space модели, или зависеть - attention механизм в трансформерах или Mamba (Selective State Spaces).

Кроме того, разные механизмы обладают разной сложностью от длины последовательности. Sequence mixing в Attention или MLP-Mixer требует квадратичного по длине последовательности числа элементарных операций с плавающей точкой (FLOPs), так как используют матричную операцию довольно общего вида. Sequence mixers, обладающие некоторой структурой (низкоранговые, Toeplitz матрицы, DFT, бабочки) позволяют добиваться субквадратичной сложности (обычно с некоторой просадкой в качестве).

И sequence mixing может быть как причинным (causal attention, большинство SSM, в частности, модная нынче Mamba 🐍), где текущий элемент последовательности может смотреть только в прошлое, и двунаправленным (как в masked language modelling, и большинстве задач с ViTами), где элементы последовательности могут изменять свое состояние, как глядя как на прошлые, так и на будущие токены.

И задача, которую, перед собой ставят авторы в данной работе - получение эффективного механизма двунаправленного sequence mixing, такого, чтобы он был с одной стороны субквадратичным (в идеале линейным по длине последовательности) и в то же время выразительным.