Muon. ինչ է լինում, երբ LLM-ի օպտիմիզատորը քաշի մատրիցան դիտարկում է որպես մատրիցա

Muon-ը նեյրոնային ցանցի քաշի մատրիցան դիտարկում է որպես մատրիցա, ոչ թե սկալյար կոորդինատների հարթ կույտ։ AdamW-ը յուրաքանչյուր տարրի համար պահում է շարժվող միջին և կոորդինատները նորմալացնում է առանձին. Muon-ը վերցնում է momentum-ի թարմացումը, դեն է նետում սինգուլյար արժեքների մեծությունները և պահում է ուղղությունները՝ այդ գործողությունը մոտարկելով մի քանի Newton–Schulz իտերացիաներով։

4096×4096 չափի transformer-ի քաշի մատրիցայում 16 միլիոնից ավելի արժեք կա. AdamW-ն տեսնում է կոորդինատներ, Muon-ը՝ իմաստալից ուղղություններ ունեցող գծային ձևափոխություն։

Muon-ը U diag(20, 3, 0.2) V^T տեսքի թարմացումը մոտավորապես վերածում է U diag(1, 1, 1) V^T-ի. Ortho(G) ≈ UV^T։ Ուղղությունները մնում են, մեծությունները հավասարվում են։

Գաղափարը հրապարակվեց 2024 թվականի հոկտեմբերին՝ NanoGPT-ի speedrun մրցույթի ընթացքում։ 2024 թվականի հոկտեմբերի 15-ին Muon-ի վրա հիմնված փորձարկումը սահմանեց ուսուցման արագության ռեկորդ՝ նախորդ արդյունքը բարելավելով մոտ 35%-ով. վաղ փորձերը նաև ցույց տվեցին, որ Q-ն, K-ն և V-ն ավելի լավ են աշխատում որպես առանձին մատրիցաներ։

SVD-ի փոխարեն՝ Newton–Schulz

Ամեն օպտիմիզացիայի քայլին լրիվ 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-ի կանոնը Muon-ի թարմացման RMS-ը համադրում է AdamW-ի հետ. նրանց փորձերը compute-օպտիմալ ուսուցման դեպքում ցույց տվեցին մոտ 2× հաշվողական արդյունավետություն՝ AdamW անալոգների ուսուցման FLOPs-ի մոտ 52%-ով։ Muon-ով ուսուցվել է նաև Moonlight-ը՝ 3B/16B MoE մոդելը, 5.7 տրիլիոն տոկենով։

Հիշողությունը նույնպես տարբերվում է. Adam-ի տիպի օպտիմիզատորները պահում են երկու moment-թենզոր, Muon-ի հիմնական վիճակը՝ մեկ momentum-բուֆեր. 100B պարամետրի դեպքում՝ 800 GB-ի փոխարեն մոտ 400 GB։ PyTorch-ն արդեն փաստագրում է torch.optim.Muon-ը մի քանի update-scaling ռեժիմով, DeepSpeed-ը՝ աջակցություն 2026 թվականի հունիսից։

AdamW-ն ամենուր փոխարինելը նպատակը չէ. Muon-ը նախատեսված է 2D թաքնված քաշի մատրիցաների համար, իսկ embeddings-ը, biases-ը, LayerNorm-ի պարամետրերն ու ելքային գլուխները մնում են AdamW-ի վրա։ 0.02 ֆիքսված learning rate-ով վաղ օրինակները կուրորեն չարժե պատճենել, քանի որ Jordan-ի, Moonshot-ի և Bernstein-ի կանոնները տարբեր են։

Վերահսկվող համեմատության մեջ ճարտարապետությունը, տվյալները, տոկենները և սարքավորումը անփոփոխ են, իսկ validation loss-ը չափվում է տոկենների, FLOPs-ի և GPU-ժամերի նկատմամբ։ Օպտիմիզատորը պարունակում է ենթադրություններ այն մասին, թե ինչ է պարամետրը. Adam-ը դրանք կոորդինատներ է համարում, Muon-ը՝ գծային օպերատորներ։