Muon: რა ხდება, როცა ოპტიმიზატორი წონის მატრიცას მატრიცად აღიქვამს

Muon ნეირონული ქსელის წონის მატრიცას მატრიცად განიხილავს და არა სკალარული კოორდინატების ბრტყელ ერთობლიობად. AdamW თითოეული ელემენტისთვის მოძრავ საშუალოს ინახავს და კოორდინატებს დამოუკიდებლად ნორმალიზებს; Muon კი momentum-ის განახლებას იღებს, სინგულარული მნიშვნელობების სიდიდეებს აგდებს და მხოლოდ მიმართულებებს ინახავს, ამ ოპერაციას კი რამდენიმე Newton–Schulz-ის იტერაციით აახლოებს. ამ იდეას dev.to-ზე გამოქვეყნებული განმარტება აღწერს.

4096×4096 ზომის transformer-ის წონის მატრიცაში 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-ის 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%-ია.

speedrun-იდან რეალურ წვრთნამდე

მასშტაბირების კითხვას Moonshot AI 2025 წლის ნაშრომში „Muon is Scalable for LLM Training“ უპასუხა. ორთოგონალიზებული მატრიცის RMS-ქცევა მის ზომებზეა დამოკიდებული, ამიტომ Moonshot-ის წესი Muon-ის განახლების RMS-ს AdamW-ისას აახლოებს; მათი ექსპერიმენტები compute-ოპტიმალურ წვრთნაზე დაახლოებით 2× გამოთვლით ეფექტურობას აჩვენებდა, შედარებადი შედეგი კი AdamW-ის ანალოგების ტრენინგის FLOPs-ის დაახლოებით 52%-ით მიიღწეოდა. Muon-ითვე იწვრთნა Moonlight — 3B/16B mixture-of-experts მოდელი — 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-ის პარამეტრები და output-თავები კი AdamW-ზე რჩება. 0.02-ის ფიქსირებული learning rate-ის მქონე ადრეული მაგალითების ბრმად კოპირება არ ღირს — Jordan-ის, Moonshot-ისა და Bernstein-ის სკალირების წესები განსხვავდება.

ექსპერიმენტში არქიტექტურა, მონაცემები, ტოკენები, ბატჩის ზომა და აპარატურა უცვლელი რჩება, იზომება კი validation loss ტოკენების, FLOPs-ისა და GPU-საათების მიხედვით. ოპტიმიზატორი იმას შეიცავს, თუ რას ნიშნავს პარამეტრი: Adam მათ კოორდინატებად აღიქვამს, Muon — წრფივ ოპერატორებად.