🤖 machine learning

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

यह शोध पत्र MONA को प्रस्तुत करता है, जो एक नवीन ऑप्टिमाइज़र है जो तीव्र स्थानीय मिनिमा (sharp local minima) से बचने के लिए Muon फ्रेमवर्क में नेस्टरोव-शैली के त्वरण (Nesterov-style acceleration) को एकीकृत करता है और 1B से 68B पैरामीटर्स तक की भाषा मॉडलों (language models) में अत्याधुनिक अभिसरण (state-of-the-art convergence) और डाउनस्ट्रीम प्रदर्शन प्राप्त करता है।

मूल लेखक: Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

प्रकाशित 2026-05-27

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप एक विशाल, सुपर-स्मार्ट रोबोट को मानवीय भाषा बोलना सिखाने की कोशिश कर रहे हैं। इसे करने के लिए, आपको उसके मस्तिष्क के भीतर लाखों छोटे-छोटे नॉब्स (knobs) को एडजस्ट करना होगा। इन नॉब्स के लिए सही सेटिंग्स खोजने की प्रक्रिया को "ट्रेनिंग" कहा जाता है, और वह टूल जिसका उपयोग आप इन नॉब्स को घुमाने के लिए करते हैं, उसे ऑप्टिमाइज़र (optimizer) कहा जाता है।

लंबे समय तक, इस काम के लिए मानक टूल AdamW था। यह भरोसेमंद है, जैसे कि एक स्थिर, सतर्क पदयात्री जो छोटे कदम उठाता है और लगातार अपने नक्शे की जांच करता रहता है। हाल ही में, Muon नामक एक नया टूल आया। Muon एक ऐसे पदयात्री की तरह है जो पूरे पर्वत श्रृंखला को एक साथ देखता है, और नॉब्स के समूहों को एक एकल इकाई के रूप में मानता है। यह Muon को तेज़ और अधिक कुशल बनाता है, लेकिन इसमें एक दोष है: यह एक "तीक्ष्ण घाटी" (sharp valley) में फंस सकता है।

समस्या: तीक्ष्ण घाटी का जाल (The Sharp Valley Trap)

कल्पना कीजिए कि रोबोट की सीखने की प्रक्रिया का परिदृश्य ऊबड़-खाबड़ पहाड़ी इलाका है।

सपाट घाटियाँ (Flat valleys): ये बसने के लिए अच्छी जगहें हैं; इसका मतलब है कि रोबोट ने एक स्थिर, सामान्य समाधान खोज लिया है जो कई स्थितियों में अच्छा काम करता है।
तीक्ष्ण घाटियाँ (Sharp valleys): ये संकरी, गहरी खाइयाँ हैं। यदि रोबोट इनमें गिर जाता है, तो उसे लग सकता है कि उसने सबसे निचला बिंदु खोज लिया है, लेकिन वास्तव में वह एक बहुत ही विशिष्ट, नाजुक स्थान पर फंस गया है। यदि आप उसे थोड़ा सा भी हिलाते हैं, तो वह अपने "परफेक्ट" स्थान से बाहर गिर जाएगा और खराब प्रदर्शन करेगा।

AdamW और Muon दोनों गलती से इन तीक्ष्ण घाटियों में गिर सकते हैं और फंस सकते हैं। उनके पास यह "महसूस" करने का कोई तरीका नहीं है कि आगे की जमीन कैसी है ताकि वे जान सकें कि कोई स्थान बहुत तीक्ष्ण है या नहीं।

समाधान: MONA (Nesterov Acceleration के साथ Muon)

इस शोध के लेखकों ने एक नया टूल बनाया जिसे MONA कहा जाता है। MONA को सोनार चश्मे (sonar goggles) वाले Muon के रूप में समझें।

यह कैसे काम करता है, इसके लिए एक सरल उपमा यहाँ दी गई है:

पुराना तरीका (Muon): Muon उस ढलान को देखता है जहाँ वह अभी खड़ा है और नीचे की ओर बढ़ता है। यह कुशलता से चलने में बहुत अच्छा है, लेकिन इसे यह नहीं पता कि आगे की जमीन एक हल्की ढलान है या किसी खड़ी चट्टान का किनारा।
नया तरीका (MONA): कदम उठाने से पहले, MONA यह देखता है कि पिछले कदम से वर्तमान कदम तक ढलान में कैसे बदलाव आया है।
- यदि ढलान अचानक बदल गई (जैसे किसी तीखी चट्टान से टकराना), तो MONA जानता है, "ओह, यह एक तीक्ष्ण घाटी है! चलिए यहाँ से बाहर निकलने के लिए ज़ोर से धक्का देते हैं।"
- यदि ढलान बहुत धीरे-धीरे बदली (एक हल्की, सपाट घाटी), तो MONA कहता है, "यह सुरक्षित लग रहा है। चलिए यहाँ बस जाते हैं।"

MONA गणित में एक विशेष "एक्सेलरेशन टर्म" (acceleration term) जोड़ता है। यह वर्तमान दिशा और पिछली दिशा के बीच के अंतर की गणना करता है। यह अंतर एक सेंसर की तरह काम करता है जो इलाके की "तीक्षक्ष्णता" का पता लगाता है। यदि इलाका तीक्ष्ण है, तो यह रोबोट को गड्ढे से बाहर धकेलता है। यदि यह सपाट है, तो यह रोबोट को आराम करने देता है।

उन्होंने क्या पाया?

शोधकर्ताओं ने तीन अलग-अलग आकार के भाषा मॉडलों (छोटे, मध्यम और विशाल, जो 1 बिलियन से 68 बिलियन पैरामीटर्स तक हैं) पर इस नए टूल का परीक्षण किया। उन्होंने इन्हें भारी मात्रा में टेक्स्ट (1 ट्रिलियन शब्दों तक) पर प्रशिक्षित किया।

बेहतर ट्रेनिंग: हर परीक्षण में, MONA ने मॉडलों को पुराने मानक (AdamW) और नए Muon दोनों की तुलना में तेज़ी से सीखने और एक बेहतर अंतिम अवस्था तक पहुँचने में मदद की।
स्मार्टर रोबोट: MONA के साथ प्रशिक्षित मॉडल सामान्य कार्यों, गणित की समस्याओं और कोडिंग में बेहतर थे। उदाहरण के लिए, एक 68-बिलियन पैरामीटर वाले मॉडल पर, MONA ने गणित और कोडिंग के बेंचमार्क पर सर्वश्रेष्ठ स्कोर प्राप्त किया।
फाइन-ट्यूनिंग (Fine-Tuning): प्रारंभिक प्रशिक्षण के बाद भी, जब उन्हें अतिरिक्त विशिष्ट प्रशिक्षण दिया गया (जैसे उन्हें विशेष रूप से कोडिंग लिखना सिखाना), तो MONA से शुरू हुए मॉडलों ने Muon से शुरू हुए मॉडलों की तुलना में बेहतर प्रदर्शन किया।

इसे व्यावहारिक बनाना (MONA-Lite)

लेखकों ने यह भी महसूस किया कि इस "सोनार" प्रणाली को जोड़ने से थोड़ा अधिक कंप्यूटर मेमोरी का उपयोग होता है। इसे ठीक करने के लिए, उन्होंने MONA-Lite नामक एक हल्का संस्करण बनाया।

उन्होंने मेमोरी डेटा को कंप्रेस करने की एक ट्रिक का उपयोग किया (जैसे हाई-डेफिनिशन वीडियो से स्टैंडर्ड डेफिनेशन में स्विच करना) और अतिरिक्त फाइलों को स्टोर किए बिना चलते समय ढलान के परिवर्तनों की गणना करने की एक विधि का उपयोग किया।
इससे आवश्यक अतिरिक्त मेमोरी लगभग 75% कम हो गई, जिससे इसे सीमित मेमोरी वाले कंप्यूटरों पर भी बिना किसी लाभ को खोए उपयोग करना आसान हो गया।

सारांश

संक्षेप में, MONA एक अपग्रेड है जो Muon ऑप्टिमाइज़र का उन्नत रूप है। यह Muon की गति और दक्षता को बनाए रखता है लेकिन इसमें एक "कर्वेचर सेंसर" (curvature sensor) जोड़ा गया है जो AI को सीखने के दौरान खराब, तीक्ष्ण स्थानों में फंसने से बचने में मदद करता है। इसके परिणामस्वरूप अधिक स्मार्ट, अधिक सक्षम भाषा मॉडल मिलते हैं जो गणित, कोडिंग और सामान्य तर्क में बेहतर होते हैं, और साथ ही इतने कुशल होते हैं कि उन्हें मानक हार्डवेयर पर चलाया जा सके।

तकनीकी सारांश: MONA – नेस्टरोव त्वरण के साथ मून ऑप्टिमाइज़र (Muon Optimizer with Nesterov Acceleration)

1. समस्या विवरण

लार्ज लैंग्वेज मॉडल्स (LLMs) का प्रशिक्षण तेजी से ऐसे ऑप्टिमाइज़र्स की मांग कर रहा है जिनमें बेहतर सैंपल एफिशिएंसी (sample efficiency) हो। जबकि एक दशक से अधिक समय से AdamW मानक रहा है, हाल ही में Muon जैसे विकल्प उभरे हैं। Muon प्रथम-क्रम विधियों (first-order methods) में सुधार करता है, जो वेट मैट्रिसेस को एकल ज्यामितीय इकाइयों (single geometric units) के रूप में मानता है और मोमेंटम बफर पर मैट्रिक्स ऑर्थोगोनलाइजेशन (न्यूटन-शुलज़ इटरेशन के माध्यम से) लागू करता है। यह दृष्टिकोण स्पेक्ट्रल नॉर्म (spectral norm) के तहत स्टाइपेस्ट डिसेंट (steepest descent) के साथ संरेखित होता है, जो बड़े पैमाने पर Mixture-of-Experts (MoE) प्रशिक्षण में AdamW की तुलना में लगभग $2\times$ कम्प्यूटेशनल दक्षता प्रदान करता है।

हालाँकि, Muon में सभी प्रथम-क्रम ग्रेडिएंट विधियों की तरह एक मौलिक सीमा है: इसमें स्थानीय ग्रेडिएंट जानकारी से परे लॉस लैंडस्केप (loss landscape) को खोजने के लिए स्पष्ट तंत्र का अभाव है। फलस्वरूप, Muon शार्प लोकल मिनिमा (sharp local minima) में फंस सकता है, विशेष रूप से बड़े-बैच प्रशिक्षण परिदृश्यों में। शार्प मिनिमा खराब सामान्यीकरण (generalization) से जुड़े होते हैं, जबकि फ्लैट मिनिमा (flatter minima) बेहतर डाउनस्ट्रीम प्रदर्शन प्रदान करते हैं। जबकि ALTO ऑप्टिमाइज़र ने शार्प मिनिमा से बचने के लिए ग्रेडिएंट अंतर के एक्सपोनेंशियल मूविंग एवरेज (EMA) पर आधारित एक त्वरण पद (acceleration term) पेश किया, लेकिन एक सामान्य अडैप्टर के रूप में इसका डिज़ाइन Muon की अत्यधिक गैर-रेखीय (nonlinear) ऑर्थोगोनलाइजेशन प्रक्रिया के साथ एकीकरण की चुनौतियाँ प्रस्तुत करता है।

2. कार्यप्रणाली: MONA

लेखक MONA (Muon Optimizer with Nesterov Acceleration) का प्रस्ताव करते हैं, जो Muon के ऑर्थोगोनलाइजेशन फ्रेमवर्क में सीधे एक कर्वेचर-अवेयर (curvature-aware) त्वरण पद को सहजता से एकीकृत करता है।

मुख्य एल्गोरिदम

MONA मानक Muon पाइपलाइन को मोमेंटम संचय और ऑर्थोगोनलाइजेशन से पहले एक त्वरण पद पेश करके संशोधित करता है। एक मैट्रिक्स पैरामीटर $W$ के लिए $k$ इटरेशन पर अपडेट की प्रक्रिया इस प्रकार है:

ग्रेडिएंट अंतर गणना: वर्तमान स्टोकेस्टिक ग्रेडिएंट $G_k$ और पिछले ग्रेडिएंट $G_{k-1}$ के बीच अंतर की गणना करें:
$D_k = G_k - G_{k-1}$
एक्सेलरेशन बफर अपडेट: इन अंतरों का एक एक्सपोनेंशियल मूविंग एवरेज बनाए रखें:
$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$
यहाँ, $\beta_a$ त्वरण की मेमोरी को नियंत्रित करता है।
कर्वेचर-अवेयर ग्रेडिएंट: कच्चे ग्रेडिएंट में स्केल्ड त्वरण पद जोड़कर एक संवर्धित ग्रेडिएंट $\tilde{G}_k$ का निर्माण करें:
$\tilde{G}_k = G_k + \alpha A_k$
गुणांक $\alpha$ त्वरण को स्केल करता है। सैद्धांतिक प्रेरणा बताती है कि $G_k - G_{k-1}$ , $H_k(\theta_k - \theta_{k-1})$ का सन्निकटन करता है, जहाँ $H_k$ हेसियन (Hessian) है। इस प्रकार, यह पद शार्प मिनिमा (जहाँ ग्रेडिएंट परिवर्तन बड़े होते हैं) से दूर की ओर संकेत करता है और फ्लैट क्षेत्रों में छोटा रहता है।
Muon ऑर्थोगोनलाइजेशन: संवर्धित ग्रेडिएंट को मानक Muon मोमेंटम और ऑर्थोगोनलाइजेशन चरणों में फीड किया जाता है:
$M_k = \mu M_{k-1} + \tilde{G}_k$
$O_k = \text{Newton-Schulz}(M_k)$
$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$

डिज़ाइन तर्क

ऑर्थोगोनलाइजेशन से पहले त्वरण पद लागू करना महत्वपूर्ण है। यदि इसे बाद में लागू किया जाता है, तो यह उस ऑर्थोगोनल संरचना को बाधित करेगा जिस पर Muon स्पेक्ट्रल-नॉर्म रेगुलराइजेशन के लिए निर्भर करता है। कच्चे ग्रेडिएंट को पहले कर्वेचर-अवेयर दिशा में बदलकर, मोमेंटम बफर इन समृद्ध दिशाओं को कैप्चर करता है, जिन्हें फिर Muon के मानक पाइपलाइन के माध्यम से प्रोसेस किया जाता है।

वेक्टर-मूल्य वाले पैरामीटर्स (जैसे एम्बेडिंग्स, बायस) के लिए, MONA मौजूदा Muon परंपराओं का पालन करते हुए AdamW का उपयोग करता है।

3. सैद्धांतिक विश्लेषण

पेपर मानक धारणाओं (L-smoothness, अनबायस्ड ग्रेडिएंट्स विद बाउंडेड वेरिएंस, और एक्सपेक्टेड डायरेक्शनल अलाइनमेंट) के तहत MONA के लिए कन्वर्जेंस विश्लेषण प्रदान करता है।

कन्वर्जेंस रेट: लेखक सिद्ध करते हैं कि MONA नॉन-कॉन्वेक्स सेटिंग्स में $O(K^{-1/2})$ की कन्वर्जेंस दर प्राप्त करता है, जो SGD और AdamW की एसिम्प्टोटिक दरों से मेल खाता है। त्वरण पद बाउंड में स्थिरांकों (constants) को प्रभावित करता है लेकिन एसिम्प्टोटिक दर को नहीं।
शार्प मिनिमम एस्केप: विश्लेषण औपचारिक रूपता से बताता है कि कैसे त्वरण पद चुनि적으로 शार्प मिनिमा का विरोध करता है। एक स्टेशनरी पॉइंट के पास, अपेक्षित त्वरण $E[A_k]$ , पिछले दिशाओं के हेसियन-वेटेड योग के समानुपाती होता है। बड़े आइजनवैल्यू (शार्प मिनिमा) वाली दिशाओं के लिए, त्वरण का परिमाण बड़ा होता है, जो ऑप्टिमाइज़र को दूर धकेलता है। छोटे आइजनवैल्यू (फ्लैट दिशाओं) के लिए, त्वरण छोटा होता है, जिससे स्थिर अभिसरण संभव होता है। यह तंत्र Muon के स्पेक्ट्रल-नॉर्म रेगुलराइजेशन को बनाए रखते हुए कर्वेचर अवेयरनेस जोड़ता है।

4. अनुभवजन्य परिणाम (Empirical Results)

MONA का मूल्यांकन तीन पैमानों के MoE प्रीट्रेनिंग मॉडल्स (1B, 6B, और 68B पैरामीटर्स) और एक सुपरवाइज्ड फाइन-ट्यूनिंग (SFT) चरण पर किया गया।

प्रीट्रेनिंग प्रदर्शन

मॉडेल्स: MOE-1B-A0d2B (400B टोकन्स), MOE-6B-A0d5B (1.2T टोकन्स), और MOE-68B-A3B (700B टोकन्स)।
मेट्रिक्स: कोड, गणितीय तर्क, सामान्य अंग्रेजी, और चीनी अकादमिक टेक्स्ट डोमेन में वैलिडेशन लॉस की निगरानी की गई।
निष्कर्ष: MONA ने तीनों मॉडल स्केल्स और सभी डोमेन में Muon और AdamW दोनों की तुलना में लगातार कम वैलिडेशन लॉस प्राप्त किया। प्रदर्शन का अंतर विशेष रूप से बड़े 68B मॉडल में स्पष्ट था, जहाँ MONA ने अधिक स्थिर कन्वर्जेंस प्रदर्शित किया।

डाउनस्ट्रीम बेंचमार्क्स (MOE-68B-A3B)

MONA के साथ प्रशिक्षित 68B मॉडल का व्यापक बेंचमार्क सूट पर मूल्यांकन किया गया:

सामान्य क्षमता: MONA ने MMLU, MMLU-Pro, CMMLU, CEVAL, और BBH सहित बेंचमार्क्स पर उच्चतम औसत स्कोर (0.4557) प्राप्त किया, जो Muon (0.4478) और AdamW (0.4382) से बेहतर है।
गणितीय तर्क और कोड जनरेशन: MONA ने GSM8K, GPQA, और विभिन्न कोड जनरेशन बेंचमार्क्स (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench) पर बेहतर प्रदर्शन दिखाया। विशेष रूप से, MONA ने उन कार्यों पर महत्वपूर्ण लाभ दिया जिनमें ट्रांसफर करने योग्य कोड रिप्रेजेंटेशन की आवश्यकता होती है।
सुपरवाइज्ड फाइन-ट्यूनिंग (SFT): जब MONA-प्रीट्रेन किया गया चेकपॉइंट कोड डेटा पर SFT से गुजरा, तो इसने Muon-प्रीट्रेन बेसलाइन की तुलना में पूरी प्रक्रिया के दौरान कम ट्रेनिंग लॉस बनाए रखा। कोड डेटा पर पोस्ट-SFT मूल्यांकन ने पुष्टि की कि MONA-इनिशियलाइज्ड मॉडल्स ने BigCode सूट के 7 में से 6 कार्यों पर Muon-इनिशियलाइज्ड मॉडल्स की तुलना में उच्च स्कोर प्राप्त किया।

दक्षता और परिनियोजन (Efficiency and Deployment)

पिछले ग्रेडिएंट ( $G_{k-1}$ ) और त्वरण बफर ( $A_k$ ) को स्टोर करने के मेमोरी ओवरहेड को संबोधित करने के लिए, लेखकों ने MONA-Lite पेश किया:

क्वांटाइजेशन: फ्लोट32 के बजाय bfloat16 (BF16) में सहायक अवस्थाओं (auxiliary states) को स्टोर करना।
स्ट्रीमिंग कंप्यूटेशन: समर्पित $G_{k-1}$ बफर की आवश्यकता को समाप्त करने के लिए इन-प्लेस ग्रेडिएंट अंतर की गणना करना।
परिणाम: इन रणनीतियों ने प्रशिक्षण की गुणवत्ता से समझौता किए बिना अतिरिक्त मेमोरी ओवरहेड को लगभग 75% कम कर दिया। कम्प्यूटेशनल ओवरहेड विश्लेषण ने दिखाया कि MONA ऑप्टिमाइज़र स्टेप स्तर पर केवल ~1% धीमा है, और एंड-टू-एंड इटरेशन समय में Muon की तुलना में कोई व्यावहारिक अंतर नहीं है।

5. महत्व और दावे

पेपर दावा करता है कि MONA सफलतापूर्वक Muon की ज्यामितीय दक्षता और कर्वेचर-अवेयर त्वरण की आवश्यकता के बीच के अंतर को पाटता है। इसका प्राथमिक महत्व है:

एल्गोरिद्मिक एकीकरण: यह प्रदर्शित करना कि त्वरण पदों को उनके ज्यामितीय गुणों को बाधित किए बिना मैट्रिक्स ऑर्थोगोनलाइजेशन फ्रेमवर्क में एकीकृत किया जा सकता है।
स्केलेबिलिटी: व्यापक रेंज के मॉडल साइज (1B से 68B पैरामीटर्स) और टोकन काउंट (1 ट्रिलियन टोकन्स तक) में विधि को मान्य करना, जो AdamW और Muon दोनों पर सुसंगत सुधार दिखाता है।
सामान्यीकरण: यह प्रमाण देना कि प्रीट्रेनिंग के दौरान शार्प मिनिमा से बचना डाउनस्ट्रीम कार्यों, विशेष रूप से जटिल तर्क और कोड जनरेशन में बेहतर सामान्यीकरण की ओर ले जाता है।
व्यावहारिकता: एक मेमोरी-कुशल संस्करण (MONA-Lite) प्रदान करना जो ऑप्टिमाइज़र को संसाधन-बाधित वातावरण के लिए व्यवहार्य बनाता है।

लेखक निष्कर्ष निकालते हैं कि MONA बड़े पैमाने पर LLM प्रशिक्षण के लिए एक स्टेट-ऑफ-द-आर्ट ऑप्टिमाइज़र है, जो Muon फ्रेमवर्क की कम्प्यूटेशनल दक्षता को बनाए रखते हुए बेहतर कन्वर्जेंस और डाउनस्ट्रीम प्रदर्शन प्राप्त करता है।