Badacze z Uniwersytetu Kalifornijskiego w Santa Barbara zaprezentowali MUON+, usprawnioną wersję algorytmu optymalizacji, która znacząco redukuje perplexity i czas uczenia dużych modeli językowych. Dzięki wprowadzeniu dodatkowego kroku normalizacji po fazie ortogonalizacji, nowa metoda pozwala na stabilniejsze trenowanie jednostek o skali od 60M do 1B parametrów.
Najważniejsze w skrócie
- MUON+ wprowadza dodatkową normalizację l2 (wierszową, kolumnową lub mieszaną) do standardowego schematu aktualizacji wag.
- Poprawa wyników na modelach GPT i LLaMA – w testach odnotowano spadek perplexity o nawet 2.02 punkty przy zachowaniu tego samego budżetu obliczeniowego.
- Większa stabilność przy długim trenowaniu (overtraining) – algorytm zachowuje przewagę nawet przy dostarczeniu 72 miliardów tokenów dla relatywnie małych modeli.
- Odporność na hiperparametry – metoda jest mniej czuła na dobór Learning Rate niż klasyczny Muon.
Ewolucja optymalizacji: Od AdamW do MUON+
W świecie głębokiego uczenia dominacja algorytmu AdamW wydaje się niezagrożona, jednak rosnące koszty trenowania modeli wymuszają poszukiwanie bardziej efektywnych rozwiązań. Jednym z najbardziej obiecujących kierunków ostatnich miesięcy stała się optymalizacja oparta na ortogonalizacji gradientów. Oryginalny optymalizator Muon, bazujący na iteracjach Newtona-Schulza, pokazał, że wymuszanie ortogonalności wag może radykalnie przyspieszyć zbieżność modelu.
Najnowsza publikacja „MUON+: Towards Better Muon via One Additional Normalization Step”, opublikowana w serwisie arXiv, idzie o krok dalej. Zespół pod kierownictwem Ruitao Zhanga z UC Santa Barbara udowodnił, że sam proces ortogonalizacji to za mało. Kluczem do odblokowania pełnego potencjału tej metody okazała się dodatkowa warstwa normalizacji aplikowana bezpośrednio po wyznaczeniu ortogonalnej macierzy aktualizacji.
Techniczne fundamenty: Dlaczego normalizacja zmienia wszystko?
Standardowy Muon działa na macierzach, a nie na wektorach, co odróżnia go od AdamW. Poprzez wymuszanie ortogonalności, zapobiega on zapadaniu się rangi (rank collapse) gradientów. MUON+ modyfikuje ten proces, wprowadzając operator:
W praktyce badacze przetestowali trzy warianty:
- Column-wise normalization: Skalowanie kolumn macierzy aktualizacji.
- Row-wise normalization: Skalowanie wierszy macierzy.
- Col-Row / Row-Col: Złożone techniki normalizacji dwukierunkowej.
Wyniki eksperymentów przeprowadzonych na architekturach GPT oraz LLaMA są jednoznaczne. Przykładowo, dla modelu GPT-Small (124M parametrów), MUON+ osiągnął perplexity na poziomie 27.64, podczas gdy bazowy Muon zatrzymał się na 29.66. To różnica, która w skali produkcyjnej przekłada się na gigantyczne oszczędności czasu pracy układów NVIDIA H100.
Stabilność w ekstremalnych warunkach
Jednym z najciekawszych aspektów badania jest zachowanie algorytmu podczas tzw. overtrainingu. Zespół wytrenował model LLaMA-350M na ogromnym zbiorze 72 miliardów tokenów z bazy FineWeb. W takim scenariuszu, gdzie stosunek liczby tokenów do parametrów wynosi aż 200:1, tradycyjne optymalizatory często wykazują niestabilność lub spowolnienie progresu. MUON+ utrzymał stabilną krzywą uczenia, kończąc proces z wynikiem 11.03 perplexity (wobec 11.48 dla bazowego Muona).
Co ważne dla inżynierów ML, MUON+ wykazuje znacznie mniejszą wrażliwość na dobór tempa uczenia (Learning Rate). W testach ablacyjnych wykazano, że nawet przy suboptymalnym, zbyt wysokim LR, spadek wydajności jest znacznie mniej drastyczny niż w przypadku starszych metod.
| Cecha | AdamW | Muon (Standard) | MUON+ |
|---|---|---|---|
| Podstawa | Adaptacyjne momenty | Ortogonalizacja | Ortogonalizacja + Normalizacja |
| Zbieżność | Standardowa | Szybka | Bardzo szybka |
| Wrażliwość na LR | Średnia | Wysoka | Niska |
| Złożoność obliczeniowa | Niska | Średnia (Newton-Schulz) | Średnia (pomijalny koszt normy) |
Warto zauważyć, że konkurencyjne podejścia, takie jak wspomniany wcześniej Muon, wymagają precyzyjnego dostrojenia, aby uniknąć eksplozji gradientów w początkowych fazach uczenia. MUON+ dzięki krokowi normalizacji l2 naturalnie ogranicza amplitudę aktualizacji wag, co czyni go bezpieczniejszym wyborem dla dużych klastrów obliczeniowych.
Dlaczego to ważne?
Wprowadzenie MUON+ to sygnał, że branża AI przesuwa punkt ciężkości z prostej skali (dodawania kolejnych warstw i miliardów parametrów) w stronę precyzyjnej inżynierii matematycznej procesu uczenia. Optymalizacja to „ukryty silnik” sztucznej inteligencji – od jego sprawności zależy, czy trenowanie modelu Llama-350M potrwa tydzień, czy trzy dni.
W kontekście obecnego deficytu mocy obliczeniowej i wysokich cen energii, każda poprawa efektywności o kilka procent jest warta miliony dolarów. MUON+ pokazuje, że drobna, zaledwie kilkuwierszowa zmiana w kodzie algorytmu (dodanie normalizacji po ortogonalizacji) może przynieść większe zyski niż kosztowna zmiana całej architektury sieci. To szczególnie istotne dla mniejszych graczy rynkowych i ośrodków akademickich, którzy dysponują mniejszą liczbą jednostek GPU. Stabilność MUON+ przy wysokich współczynnikach token-to-parameter sugeruje również, że metoda ta będzie idealna do tworzenia ekstremalnie wydajnych „małych” modeli (SLM), które swoimi możliwościami mogą dorównać znacznie większym jednostkom.
Co dalej?
- Skalowanie do modeli 70B+: Kolejnym logicznym krokiem będzie sprawdzenie, czy korzyści z MUON+ utrzymują się przy trenowaniu gigantycznych modeli klasy DeepSeek-V3 lub nowych wersji GPT.
- Integracja z bibliotekami optymalizacyjnymi: Można oczekiwać, że implementacje MUON+ szybko trafią do popularnych frameworków, takich jak PyTorch czy Hugging Face
accelerate, umożliwiając społeczności open-source szybsze trenowanie własnych modeli. - Badania nad dynamiczną normalizacją: Badacze sugerują, że kierunek normalizacji (wierszowa vs kolumnowa) może być dobierany dynamicznie w trakcie uczenia, co mogłoby jeszcze bardziej podnieść efektywność algorytmu.
Źródła
- arXiv – MUON+: Towards Better Muon via One Additional Normalization Step – https://arxiv.org/html/2602.21545v1
- ArXiv – Publikacja bazowa o optymalizacji ortogonalnej (Jordan et al.) – https://arxiv.org/abs/2410.10813
