Tilde Research представила Aurora — оптимизатор, устраняющий «смерть нейронов» в Muon

Tilde Research выпустила Aurora, новый оптимизатор, который устраняет системную проблему Muon: массовую «смерть» нейронов в высоких (tall) матрицах MLP. Релиз включает открытый код, 1.

Tilde Research анонсировала Aurora — оптимизатор, призванный устранить системную ошибку в Muon, из‑за которой при обучении выключается значительная доля нейронов в MLP‑слоях и они становятся «перманентно мёртвыми». Это важно для крупных тренировок: массовая потеря активных нейронов снижает эффективность использования параметров и ухудшает распространение градиента в tall‑матрицах, что влияет на качество и стабильность обучения. Muon обновляет веса через ортогонализованный градиент: при SVD G = UΣVᵀ берётся polar(G) = UVᵀ, и затем W ← W − η·UVᵀ. Практическая реализация опирается на матричные умножения и итерационные алгоритмы для вычисления полярного фактора, что сделало Muon привлекательным в соревнованиях по скорости обучения и при масштабных запусках на фронтир‑уровне.

Авторы заметили, что в tall‑матрицах, характерных для SwiGLU‑на основе MLP, Muon вводит сильную анизотропию норм строк: нельзя одновременно сохранить строгую ортогональность столбцов и обеспечить равномерные по строкам обновления. В результате некоторые нейроны получают чрезмерно большие апдейты, а другие почти не получают сигналов; в экспериментах к 500‑му шагу тренировки более одной из четырёх нейронов становились практически неактивными, и этот дефицит сигнала распространялся на последующие слои.

Промежуточный фикс — U‑NorMuon — вырос из идеи NorMuon, который нормализует каждую строку до единичной нормы. Для column‑orthogonal tall матрицы правильная средняя норма строки равна √(n/m), а не 1, поэтому U‑NorMuon нормализует строки до √(n/m). В тестах на модели масштаба 340M U‑NorMuon превосходит Muon и NorMuon, полностью устраняя феномен «смерти» нейронов и делая leverage‑оценки изотропными, но при этом он жертвует точностью полярного фактора.

Жертва точности полярного фактора в U‑NorMuon теоретически нежелательна и эмпирически дорого обходится: более точная ортогонализация в стандартном Muon даёт более низкие значения потерь. Aurora решает этот компромисс, формулируя задачу как «steepest descent under two joint constraints» — метод обеспечивает равномерность апдейтов по нейронам без потери точности полярного фактора, сохраняя преимущества ортогонализации. Практические следствия заметны в эксперименте: релиз сопровождается открытым кодом, 1.1‑миллиардным экспериментом предобучения и результатом state‑of‑the‑art на модифицированном nanoGPT speedrun. Исправление скрытой потери активных нейронов стабилизирует поток градиентов через MLP‑проекции, улучшает использование параметров в tall‑матрицах и даёт воспроизводимые выгоды при предобучении и скоростных проверках.

Публикация с кодом даёт командам, использующим Muon в масштабных тренировках, возможность воспроизвести результаты и при необходимости перейти на Aurora или применить промежуточные правки вроде U‑NorMuon, чтобы уменьшить анизотропию и сохранить эффективность ортогонализации.