📌MesaNet: оптимальная адаптация весов в реальном времени.

MesaNet — это новое поколение RNN-архитектур от команды Google Paradigms of Intelligence, созданное для эффективного моделирования длинных последовательностей (например, текста или временных рядов).

Ключевая задача MesaNet - преодолеть главный недостаток трансформеров: линейный рост вычислительных затрат и потребления памяти при увеличении длины последовательности во время инференса. В отличие от классических подходов, MesaNet достигает постоянной сложности на токен за счёт «оптимального обучения в реальном времени» — динамической подстройки внутренних весов под контекст прямо во время работы модели, без хранения всей истории токенов.

Архитектурно, MesaNet построен как стек чередующихся блоков: Mesa-слои (для смешивания информации вдоль последовательности) и MLP (для обработки признаков внутри токена).

Mesa-слой - это сердце системы. Вместо стандартного обновления весов через градиентный спуск (как в Mamba или DeltaNet), он решает оптимизационную задачу для каждого нового токена: ищет матрицу весов, минимизирующую квадратичную ошибку предсказания на всей текущей последовательности.

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

Еще одна, не менее важная опция — динамическое распределение вычислений. Число шагов сопряженного градиента не фиксировано: сложные последовательности требуют больше итераций для сходимости. Это позволяет гибко балансировать точность и скорость.

Сравнение с трансформерами (MHA) и современными RNN (Mamba2, xLSTM, DeltaNet) на синтетике (MAD, RegBench) и языке (SlimPajama) показало: MesaNet сопоставим с трансформерами по perplexity, но выигрывает у других RNN на задачах, требующих длинного контекста. При этом он сохраняет преимущество RNN — постоянные память/вычисления на токен при инференсе.

Интересный паттерн выявили во время тестов: MesaNet, да и просто RNN, точнее предсказывают ранние токены последовательности, а трансформеры - поздние. На длинной экстраполяции (до 32k токенов) MesaNet обошла Mamba2 и xLSTM, но уступила трансформеру.


🔜 Посмотреть видео с докладом про работу.


@ai_machinelearning_big_data

#AI #ML #RNN #MesaNet