Новый субквадратичный оператор HyenaND обходит attention в задачах с многомерными данными

Механизм внимания (attention) стал стандартом в задачах обработки последовательностей, но его квадратичная сложность ограничивает применение к многомерным данным, таким как изображения, объёмы и решения дифференциальных уравнений. Субквадратичные альтернативы, например классические свёртки, лишены глобального контекста, а рекуррентные модели вынуждены «сплющивать» данные в одномерную последовательность, разрушая пространственную структуру.

В новой работе на arXiv представлен HyenaND — оператор, который обходит эти ограничения. Он использует многомерные свёртки со скрытой параметризацией ядра, зависящей от входных данных. Благодаря этому HyenaND сохраняет глобальное поле восприятия и адаптируется к конкретному сигналу, оставаясь при этом субквадратичным.

Авторы реализовали оператор в библиотеке nSubQ на CUDA. В ней этап свёртки через быстрое преобразование Фурье (FFT) оптимизирован так, что теоретические O(L log L) превращаются в реальный прирост скорости на GPU. Это делает подход практичным для больших многомерных массивов.

Эксперименты проведены в четырёх областях: геномика (длинные последовательности), компьютерное зрение, медицинская визуализация и моделирование с помощью PDE. Чистые стеки HyenaND достигают точности сильных baselines attention, а гибридные конфигурации с чередованием HyenaND и attention превосходят как чистый attention, так и гибриды на основе рекуррентных блоков.

Работа показывает, что можно получить глобальный контекст и входную зависимость без квадратичного роста вычислений. Это открывает путь к созданию эффективных моделей для работы с 2D, 3D и более высокими размерностями, где attention дорог, а статические свёртки недостаточно гибки.

HyenaND может найти применение в анализе медицинских снимков, обработке 3D-сцен, климатическом моделировании и других областях, где важна как точность, так и производительность. Разработка находится на стадии исследования, но код уже доступен.