Ортогонализация матриц улучшает память в рекуррентных моделях
Трансформеры хорошо решают задачи ассоциативного припоминания (AR), механизм внимания даёт каждому токену прямой доступ ко всем предыдущим токенам. Но в некоторых областях (например, длинногоризонтное обучение с подкреплением в стиле Dreamer) квадратичные затраты трансформеров неприемлемы. Исследователи ищут способы улучшить рекуррентные сети (RNN), сохранить их эффективность, но вернуть способность к припоминанию.
Самый успешный RNN для AR, mLSTM, вариант LSTM, который поддерживает матричную память. На тесте MQAR mLSTM показывает лучшие результаты, чем базовые модели, но это может быть недостаточно: в реальных условиях (когда переходы в окружении зашумлены) нужен более тонкий тест, noisy associative recall (NAR). Задача NAR: дана последовательность пар ключ-значение, перемешанная с отвлекающими токенами (дистракторами); модель должна предсказать нужное значение, имея в голове много ключей и много шума.
Подсказку дал оптимизатор Muon, который добился успеха в языковом моделировании благодаря ортогонализации импульсов (momentum): он уравнивает представляемые направления, не давая нескольким сильным направлениям задавить слабые. Свежие исследования показали, что Muon особенно хорош для хвостовых (tail-end) задач ассоциативной памяти. Авторы решили проверить: улучшит ли ортогонализация матричной памяти mLSTM во время считывания показатели NAR.
В экспериментах на MAD noisy-recall тесте (с 80%-ным уровнем шума, разные размеры словаря и длины последовательностей) ортогонализированный вариант mLSTM бил базовый mLSTM по скорости успеха и средней точности. Поразительнее всего: разница растёт с усложнением (в режиме vocab-96 базовый mLSTM едва решал задачу, 4 из 24 обученных моделей, ортогонализированный, 14, 16). Улучшение достигалось на фиксированном количестве параметров: ортогонализация через пять итераций метода Ньютона, Шульца добавляла дополнительный выигрыш ценой больше вычислений.
Однако авторы осторожны в интерпретации: результаты получены на малых моделях в синтетическом сеттинге NAR. Нужно проверить, переносятся ли эти выигрыши на реальные бенчмарки для больших моделей. Работа финансирована Paradigm.
Ключевые факты
- Ортогонализация матричной памяти mLSTM (рекуррентной сети) улучшает задачи ассоциативного припоминания в зашумленных условиях (NAR), критичной для длинногоризонтных RL-приложений
- Метод позаимствован из оптимизатора Muon: ортогонализация через нормализацию по норме Фробениуса + пять итераций Ньютона, Шульца, уравнивает представляемые направления, не дав сильным доминировать
- На синтетическом тесте MAD (80% шума) ортогонализированный mLSTM вытягивает базовый из состояния почти отказа (4/24 решений) к надёжной работе (14, 16/24), особенно на сложных конфигурациях
- Улучшение достигается на фиксированном наборе параметров; дополнительные итерации Ньютона, Шульца добавляют выигрыш, компромисс, больше FLOPs и wall-clock времени
- Авторы подчёркивают: результаты на малых моделях в синтетическом сеттинге; переход на реальные бенчмарки больших моделей требует отдельной проверки
Почему это важно
Трансформеры доминируют благодаря механизму внимания, но в областях с длинными горизонтами или критичным потреблением памяти их квадратичная сложность неприемлема. RNN давно хотели вернуть в строй: они линейны, но историически уступают по ассоциативному припоминанию. Показанный метод ортогонализации матричной памяти, конкретный шаг к RNN, которая не теряет памяти.
Кому это важно
Разработчикам систем длинногоризонтного обучения с подкреплением (RL agents, планирование в больших окружениях), инженерам, оптимизирующим модели под ограничения по памяти и задержке, исследователям в архитектурах нейросетей, особенно тем, кто ищет альтернативы трансформерам.
Как это применить
Метод не требует внедрения новой архитектуры, добавляются шаги ортогонализации в процедуру считывания памяти mLSTM: нормализация по норме Фробениуса и 5 итераций Ньютона, Шульца. Реализация откомпилирована и опубликована. Можно встроить в существующие mLSTM-базированные системы; компромисс, небольшое замедление вычислений.
Можно ли доверять
Исследование финансировано Paradigm, авторы указывают на собственные ограничения (малые модели, синтетический тест). Метод опирается на проверенный в промышленности оптимизатор Muon и теорию ортогональных матриц. Но результаты пока не проверены на реальных приложениях больших моделей, это честно заявлено.
Риски и подводные камни
Основной риск, синтетичность тестов (NAR): неясно, улучшит ли ортогонализация реальные задачи RL на полноразмерных моделях. Вычислительные затраты (5 итераций Ньютона, Шульца) могут перевесить прирост в edge-случаях. Авторы не интегрировали метод в реальные системы Dreamer-типа, это стоит на повестке.