Трансформеры отлично справляются с ассоциативным recall — внимание даёт каждому токену прямой доступ к предыдущим. Но для некоторых задач (например, long-horizon RL в стиле Dreamer) квадратичная сложность внимания непозволительна. Тут нужны рекуррентные сети, и лучшая из них для recall — mLSTM, вариант LSTM с матричной памятью. На бенчмарке MQAR mLSTM показывает улучшенный recall, но он не измеряет шумный recall (NAR). А в задачах с шумными переходами среды NAR — более адекватный тест.
Исследователи взяли синтетическую задачу MAD noisy recall, где ключи, значения и дистракторы используют разные диапазоны токенов. И поставили вопрос: как улучшить NAR у mLSTM? Ответ подсмотрели у оптимизатора Muon, который ортогонализирует моменты, не давая сильным направлениям доминировать и выравнивая представления. Недавно показали, что Muon обгоняет Adam в изучении хвостовых ассоциативных воспоминаний — идея в том, что ортогонализация не даёт слабым воспоминаниям вытесняться.
Авторы решили ортогонализировать матрицу памяти mLSTM во время чтения (но не записывать её обратно — это ухудшало результат) и тренироваться с этим дополнительным процессом. Использовали нормализацию по Фробениусу и пять итераций Newton-Schulz, позволив градиентам протекать через процедуру.
Сравнили ортогонализированные mLSTM с базовыми на предсказании следующего токена в задаче MAD noisy AR. Все модели обучались с AdamW (betas=0.9/0.999, weight_decay=0.01) 2000 шагов, batch size 64, learning rate подбирали перебором. Для каждой комбинации размера словаря (80 или 96 токенов) и длины последовательности (512, 768, 1024) запускали по 24 seed'а с разными LR.
Результаты: ортогонализация улучшила и процент успешных seed'ов (выше 80% точности), и среднюю точность по всем конфигурациям. Прирост особенно заметен на сложных режимах (vocab 96, длина 768/1024): базовая mLSTM почти проваливалась (4 из 24 seed'ов решены), а ортогонализированная выдавала 14–16 решённых seed'ов. Например, на vocab 80 / длина 1024 точность выросла с 83.3% до 98.5%, а на vocab 96 / длина 1024 — с 23.1% до 68.5%. Авторы называют это впечатляющим для «маленького вмешательства» — Newton-Schulz даёт дополнительные gains при фиксированном числе параметров, хотя и требует больше FLOP'ов и времени.
Впрочем, они предупреждают: результаты получены на маленьких моделях и синтетической задаче. Стоит проверить, переносятся ли выигрыши в NAR на реальные бенчмарки с большими моделями.