Muon: що станеться, коли оптимізатор LLM сприймає вагову матрицю як матрицю

Muon розглядає вагову матрицю нейронної мережі саме як матрицю, а не як плаский набір скалярних координат. AdamW зберігає ковзне середнє для кожного елемента й нормалізує координати незалежно; Muon бере momentum-оновлення, відкидає величини сингулярних значень і зберігає сингулярні напрямки, наближаючи цю операцію кількома ітераціями Newton–Schulz. Цю ідею простежує пояснення на dev.to.

Вагова матриця трансформера 4096×4096 містить понад 16 мільйонів значень: AdamW бачить 16 мільйонів координат, Muon — лінійне перетворення зі змістовними напрямками.

Muon перетворює оновлення, наприклад U diag(20, 3, 0.2) V^T, приблизно на U diag(1, 1, 1) V^T: Ortho(G) ≈ UV^T. Напрямки залишаються, величини вирівнюються.

Ідея публічно з'явилася у жовтні 2024 року під час змагання зі швидкісного тренування NanoGPT. 15 жовтня 2024 року запуск на базі Muon встановив рекорд швидкості навчання, покращивши попередній результат приблизно на 35%; тогочасні експерименти також показали, що Q, K і V краще працюють як окремі матриці.

Newton–Schulz замість SVD

Повний SVD на кожному кроці оптимізатора був би надто повільним, тому Muon його ніколи не обчислює. Він ділить матрицю на її норму Frobenius і багаторазово застосовує поліном — X_next = aX + b(XX^T)X + c(XX^T)²X — з підібраними коефіцієнтами, зазвичай за п'ять ітерацій. Оскільки поліном діє лише на сингулярні значення, вони прямують до 1, а сингулярні вектори залишаються на місці; уся операція ефективно працює в bfloat16, а накладні витрати становлять близько 0,5% у прикладі автора.

Від спідранів до реальних тренувань

Питання масштабування розглянула Moonshot AI у статті 2025 року «Muon is Scalable for LLM Training». Поведінка RMS ортогоналізованої матриці залежить від її розмірів, тому правило Moonshot робить RMS оновлення Muon порівнянним з AdamW; їхні експерименти показали приблизно 2× обчислювальну ефективність за compute-оптимального навчання та порівнянну якість приблизно за 52% FLOPs від аналогів на AdamW. На Muon також навчили Moonlight — модель mixture-of-experts 3B/16B — на 5,7 трильйона токенів.

Пам'ять теж відрізняється: оптимізатори типу Adam зберігають по два moment-тензори на параметр, а основний стан Muon — один momentum-буфер: за 100B параметрів це близько 400 GB замість 800 GB. PyTorch уже документує torch.optim.Muon з кількома режимами масштабування оновлень, а DeepSpeed додав підтримку в червні 2026 року.

Замінити AdamW всюди — не мета: Muon призначений для 2D прихованих вагових матриць, тоді як embeddings, biases, параметри LayerNorm і вихідні голови залишаються на AdamW. Ранні приклади з фіксованим learning rate 0.02 не варто копіювати наосліп, бо масштабування Jordan, Moonshot і Bernstein різняться.

У контрольованому порівнянні архітектура, дані, токени, розмір батчу й обладнання залишаються незмінними, а validation loss вимірюється щодо токенів, FLOPs і GPU-годин. Оптимізатор містить припущення про те, чим є параметр: Adam вважає їх координатами, Muon — лінійними операторами.