Muon: מה קורה כשמאמץ LLM מתייחס למטריצת משקלים כמו למטריצה

Muon מתייחס למטריצת המשקלים של רשת נוירונים כמו למטריצה, לא כשקית שטוחה של קואורדינטות סקלריות. AdamW שומר ממוצע נע לכל רכיב ומנרמל כל קואורדינטה בנפרד; Muon לוקח את עדכון ה-momentum, משליך את גודלי הערכים הסינגולריים ומשאיר את הכיוונים, תוך קירוב הפעולה בכמה איטרציות Newton–Schulz. הסבר ב-dev.to מאת Shrijith Venkatramana עוקב אחרי הדרך של הרעיון.

מטריצת משקלים בגודל 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, בזמן תחרות ה-speedrun של 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× באימון אופטימלי-חישובית ועל ביצועים דומים בכ-52% מה-FLOPs של מקבילי AdamW. הם גם אימנו את Moonlight, מודל mixture-of-experts בגודל 3B/16B, על 5.7 טריליון טוקנים.

גם הזיכרון שונה: אופטימייזרים מסוג Adam שומרים שני טנסורי moment לכל פרמטר, ואילו המצב המרכזי של Muon כולל מאגר momentum אחד — כ-400 GB במקום 800 GB ב-100B פרמטרים. PyTorch כבר מתעד torch.optim.Muon עם כמה מצבי סקייל, ו-DeepSpeed הוסיפה תמיכה ביוני 2026.

החלפה של AdamW בכל מקום אינה הכוונה: Muon מיועד למטריצות משקל נסתרות דו-ממדיות, בעוד embeddings, biases, פרמטרי LayerNorm וראשי פלט נשארים על AdamW. דוגמאות מוקדמות עם learning rate קבוע של 0.02 לא כדאי להעתיק בעיניים עצומות, מאחר שהסקייל של Jordan, Moonshot ו-Bernstein שונה.

בהשוואה מבוקרת ארכיטקטורה, נתונים, טוקנים, גודל באצ' וחומרה נשארים קבועים, ו-validation loss נמדד מול טוקנים, FLOPs ושעות GPU. האופטימייזר מכיל הנחות לגבי מהו פרמטר: Adam מתייחס לפרמטרים כקואורדינטות, Muon — כאופרטורים לינאריים.