Po dekádě Adamu: čím se neuronové sítě doopravdy trénují rychleji

Qminers Reading Group, session 13. Tento článek kopíruje formát našich interních setkání: co říkají články a co si o nich myslíme my.

Optimalizace je tichý motor hlubokého učení. Každý model, od malé prognostické sítě po špičkový LLM, vzniká tak, že se optimalizační algoritmus prokousává krajinou chybové funkce. Zhruba deset let byl tím algoritmem Adam a otázka se zdála uzavřená. Éra velkých jazykových modelů ji ale otevřela znovu: když jeden trénink stojí miliony, znamená desetiprocentní úspora reálné peníze, a o tu se uchází celá vlna nových metod. Session 13 prošla vývojovou linii od klasického stochastického gradientního sestupu až po dnešní špičku, celkem sedm článků a jeden vlivný blogový příspěvek.

Co říkají články

Klasika. Stochastický gradientní sestup (SGD) posouvá parametry proti směru gradientu, vždy podle jedné zašuměné dávky dat. Momentum tento šum vyhlazuje tím, že průměruje nedávné gradienty. AdaGrad (2011) přinesl myšlenku, ze které vychází všechno pozdější: velikost kroku každého parametru se škáluje podle historie jeho gradientů, takže parametry, které se aktualizují jen zřídka, dělají větší kroky. Adam (2014) spojil momentum s adagradovským škálováním a přidal korekci vychýlení. Jeho výchozí nastavení se ukázalo tak spolehlivé, že se ho většina praktiků nikdy ani nedotkla. AdamW (2017) opravil jemnou chybu v tom, jak Adam zachází s regularizací: u adaptivních metod není L2 regularizace a útlum vah (weight decay) totéž, a jejich oddělení umožňuje ladit rychlost učení a sílu útlumu nezávisle na sobě. AdamW zůstává výchozí volbou pro trénink jazykových modelů dodnes.

Nová vlna. Adam s každým parametrem zachází jako se samostatným číslem. Jenže parametry neuronové sítě mají podobu matic a gradient matice v sobě nese strukturu, kterou čistě skalární metody zahazují. Shampoo (2018) toho využívá: každou gradientní matici předpodmiňuje zleva i zprava a plnou informaci druhého řádu přitom aproximuje Kroneckerovým součinem, čímž nezvládnutelnou inverzi matice 50 000 × 50 000 promění v inverzi dvou malých matic. Článek roky ležel bez povšimnutí, dokud ho distribuovaná implementace v PyTorchi (2023) neučinila prakticky použitelným. Shampoo pak vyhrálo hlavní soutěžní kategorii benchmarku AlgoPerf, dosud nejseriózněji vedeného nezávislého srovnání tréninkových algoritmů.

SOAP (2024) ukázal, že krok algoritmu Shampoo odpovídá spuštění Adafactoru ve vlastní bázi Shampoo, a celý postup zjednodušil na jednu větu: spusťte Adam, ale v pootočené souřadné soustavě, kterou obnovujete jen občas. Výsledek dosahuje stejných zisků jako Shampoo s menším počtem laditelných parametrů a na modelu s 660 miliony parametrů uvádí zhruba o 40 % méně iterací než pečlivě naladěný AdamW. Muon volí ještě radikálnější zkratku: vezme matici momenta, ortogonalizuje ji (rychlou Newtonovou-Schulzovou iterací místo nákladného singulárního rozkladu) a udělá krok tímto směrem. Z aktualizace se tak stane čistá rotace v maticovém prostoru, kde jsou všechny singulární hodnoty stlačeny na jedničku. Muon spotřebuje oproti Adamu poloviční paměť a používá se jen na dvourozměrné váhy skrytých vrstev, zatímco embeddingy a normalizační vrstvy má stále na starosti AdamW.

Konfrontace s realitou. Článek s naším nejoblíbenějším názvem, Fantastic Pretraining Optimizers and Where to Find Them, poměřuje celý ten zvěřinec proti poctivě naladěnému AdamW. Vyčnívají tři zjištění. Zrychlení nepřekračuje 1,4násobek, což je hluboko pod dvojnásobky z původních článků, které většinou jen odrážejí slabě nastavené srovnání. S rostoucí velikostí modelu zisky slábnou: maticové metody přinášejí zhruba 1,3násobek u modelů do 520 milionů parametrů, ale jen asi 1,1násobek u modelů s 1,2 miliardy parametrů trénovaných na osminásobku chinchillového poměru. A rané průběhy chybové funkce klamou: pořadí optimalizátorů se během poklesu rychlosti učení může přeházet, takže posuzovat metody podle průběžných kontrolních bodů je nespolehlivé.

Rozvrhy učení. Článek The Road Less Scheduled zpochybňuje jiný zažitý předpoklad. Teorie říká, že optimální rychlost učení je konstantní a že výsledkem má být průměr všech iterací; praxe používá klesající rozvrhy a ponechává si poslední iteraci. Autoři ukazují, že jde o dva pohledy na tentýž trik: lineární pokles váží gradienty stejně, jako by to dělalo průměrování iterací, a všudypřítomný kosinový rozvrh funguje hlavně proto, že lineární pokles přibližuje, přičemž posledních pár procent běhu promrhá na rychlostech učení příliš malých, než aby na nich záleželo. Jejich metoda bez rozvrhu (schedule-free) dosahuje stejně optimálních rychlostí, aniž by dopředu potřebovala znát celkový počet kroků, a v samoladicí kategorii AlgoPerf skončila na prvním místě.

Co si myslíme my

Nejcennější z celé osmičky je ten demytizující článek. Ne proto, že by nové optimalizátory byly k ničemu, ale proto, že vyčísluje, jakou mají skutečnou hodnotu: dobře naladěný AdamW je mnohem tvrdší soupeř, než literatura naznačuje, a nastavení se mezi optimalizátory nepřenáší, takže většina publikovaných srovnání je nefér už ze své podstaty.

I tak platí, že 1,1 až 1,4násobek není nic zanedbatelného. Při rozpočtech na trénink velkých jazykových modelů jde o velké množství výpočetního výkonu, a pro týmy, které často přetrénovávají středně velké modely (přesně v tomto režimu maticové metody vynikají nejvíc), je důvod k odchodu od Adamu teď věrohodný. Zjištění, že optimální volba se posouvá podle poměru dat k velikosti modelu (Muon vítězí kolem jednonásobku chinchilly, Kron a SOAP přebírají vedení od osminásobku výš), je přesně ten druh praktického detailu, který odlišuje benchmark od marketingu.

Výsledek metody bez rozvrhu změnil některým z nás pohled na rozvrhy učení jako takové. Tvrzení, že kosinový pokles, použitý v tisících článků, je jen hrubší aproximací lineárního poklesu, který je sám náhradou za průměrování iterací, je vzácný případ, kdy teorie čistě vysvětlí zažitou empirickou praxi. Náladu skupiny shrnul jeden slajd z přednášky autorů: přestaňte používat kosinus.

Náš celkový závěr: po deseti letech je otázka optimalizátoru znovu otevřená. Odpověď už není „prostě použij Adam", ale není jí ani žádný jediný nástupce. Záleží na velikosti modelu, poměru dat a na tom, kolik máte prostoru na ladění. Což je, jakkoli to zní nepohodlně, poctivá odpověď na většinu otázek v tomto oboru.

Doporučené pořadí čtení: články o Adamu a AdamW jako základ, články o Shampoo a SOAP kvůli linii maticového předpodmiňování, text o Muonu jako současný favorit a Fantastic Pretraining Optimizers dřív, než uvěříte jakémukoli tvrzení o zrychlení.