तकनीकी सारांश: Muonp – भिन्नात्मक स्पेक्ट्रल शक्तियों के साथ Muon

समस्या विवरण (Problem Statement)

Muon ऑप्टिमाइज़र डीप न्यूरल नेटवर्क, विशेष रूप से लार्ज लैंग्वेज मॉडल्स (LLMs) को प्रशिक्षित करने के लिए एक शक्तिशाली उपकरण के रूप में उभरा है। यह ग्रेडिएंट मैट्रिक्स $G = USV^\top$ को इसके पोलर फैक्टर $UV^\top$ से बदलकर काम करता है, जिससे सिंगुलर वैल्यू स्पेक्ट्रम को समतल (flattening) किया जाता है। यह "स्पेक्ट्रम फ्लैटनिंग" विभिन्न सिंगुलर दिशाओं में संतुलित शिक्षण को बढ़ावा देती है और प्रमुख मोड को प्रशिक्षण अपडेट पर एकाधिकार करने से रोकती है।

हालाँकि, लेखक Muon के दृष्टिकोण में एक महत्वपूर्ण सीमा की पहचान करते हैं: पूर्ण फ्लैटनिंग सभी सिंगुलर-वैल्यू जानकारी को त्याग देती है। जबकि यह एकरूपता प्रीट्रेनिंग के दौरान व्यापक फीचर डिस्कवरी के लिए फायदेमंद है, यह अनुकूलन कार्यों (जैसे फाइनट्यूनिंग) के लिए इष्टतम नहीं हो सकती है जहाँ सिंगुलर मानों का परिमाण (magnitude) यह बताता है कि कौन सी अवशिष्ट दिशाएँ (residual directions) बदलने के लिए सबसे अधिक महत्वपूर्ण हैं। यह कार्य जिस केंद्रीय प्रश्न को संबोधित करता है, वह है: क्या हम Muon के मैट्रिक्स-स्तरीय स्पेक्ट्रल लाभों को बनाए रखते हुए, प्रभावी अनुकूलन के लिए आवश्यक परिमाण संरचना को सुरक्षित रख सकते हैं?

कार्यप्रणाली (Methodology)

सैद्धांतिक सूत्रीकरण: Muonp

लेखक Muonp पेश करते हैं, जो Muon ऑप्टिमाइज़र का एक सामान्यीकरण है जो भिन्नात्मक स्पेक्ट्रल शक्तियों (fractional spectral powers) का उपयोग करके वेट्स को अपडेट करता है:
$$\text{Update} \propto US^pV^\top$$
जहाँ $p \in (0, 1)$ एक परिमेय संख्या (rational number) है।

• ज्यामितीय व्याख्या: लेखक सिद्ध करते हैं (Theorem 2.1) कि यह अपडेट नियम Schatten $q$-norm के तहत सबसे तीव्र ढाल (steepest-descent) की दिशा के अनुरूप है, जहाँ $q = 1 + 1/p$।
• इंटरपोलेशन (Interpolation):
  • $p = 0$: मानक Muon ($UV^\top$) के अनुरूप है, जो अधिकतम एकरूपता को अधिकतम करता है।
  • $p = 1$: मानक ग्रेडिएंट डिसेंट ($USV^\top = G$) के अनुरूप है।
  • $0 < p < 1$: इन दोनों के बीच इंटरपोलेट करता है, आंशिक परिमाण जानकारी को बनाए रखता है।

एल्गोरिद्मिक चुनौती और समाधान

$US^pV^\top$ को सीधे कंप्यूट करने के लिए सिंगुलर वैल्यू डीकंपोजिशन (SVD) की आवश्यकता होती है, जो बड़े मॉडल्स के लिए अत्यधिक गणनात्मक रूप से महंगा है। लेखक बिना स्पष्ट SVD के इस अपडेट को अनुमानित करने की चुनौती को हल करते हैं:

1. यूनिवेरिएट पुनरावृत्ति की असंभवता: लेखक सिद्ध करते हैं (Lemma 2.2) कि कोई भी निश्चित यूनिवेरिएट बहुपद पुनरावृत्ति (univariate polynomial iteration) $S_{n+1} = f(S_n)$ किसी भी परिमेय $p$ के लिए $S^p$ की ओर अभिसरित (converge) नहीं हो सकती। पोलर फैक्टर ($p=0$) के विपरीत, भिन्नात्मक शक्तियों के लिए एल्गोरिदम को मूल मैट्रिक्स को "याद रखने" की आवश्यकता होती है।
2. बाइवेरिएट बहुपद पुनरावृत्ति (Bivariate Polynomial Recurrence): इसे हल करने के लिए, लेखक निम्न-डिग्री विषम द्विचरीय पुनरावृत्तियों (low-degree odd bivariate recurrences) को व्युत्पन्न करते हैं। एल्गोरिदम पुनरावर्ती रूप से मैट्रिसेस $G_n = US_nV^\top$ की गणना करता है जहाँ सिंगुलर मान $S_{n+1} = f(S_n, S)$ के माध्यम से विकसित होते हैं।
   • Theorem 2.3 प्रदर्शित करता है कि किसी भी विषम द्विचरीय बहुपद $f(x, y)$ के लिए, पद $Uf(S_n, S)V^\top$ को केवल मैट्रिक्स गुणन और ट्रांसपोज़ का उपयोग करके कंप्यूट किया जा सकता है, जो Muon की गणनात्मक संरचना को सुरक्षित रखता है।
   • Theorem 2.4 स्थापित करता है कि ऐसे बहुपदों का अस्तित्व है जो किसी भी परिमेय $p \in (0, 1)$ के लिए $S^p$ की ओर अभिसरित होते हैं।
3. स्पष्ट निर्माण: पेपर इन बहुपदों के लिए स्पष्ट सूत्र प्रदान करता है (जैसे, $p=1/3$ के लिए $f(x, y) = x + c(y - x^3)$)। गुणांकों को यह सुनिश्चित करने के लिए चुना गया है कि $S^p$ एक आकर्षक फिक्स्ड पॉइंट (attracting fixed point) हो।

कार्यान्वयन विवरण (Implementation Details)

• जटिलता: इसकी गणनात्मक जटिलता $O(\omega k N)$ बनी रहती है, जहाँ $\omega$ मैट्रिक्स गुणन जटिलता है, $k$ प्रति चरण गुणनों की संख्या है (बहुपद डिग्री पर निर्भर), और $N$ पुनरावृत्तियों की संख्या है। Muon की तुलना में Muonp केवल मैट्रिक्स गुणनों की संख्या में एक स्थिरांक कारक (constant factor) से भिन्न है।
• स्थिरता: इसे bfloat16 में लागू करने पर यह संख्यात्मक रूप से स्थिर है, जो इसे आधुनिक हार्डवेयर के लिए उपयुक्त बनाता है।
• करिकुलम संगतता (Curriculum Compatibility): चूंकि Muon ($p=0$) Muonp का एक विशेष मामला है, इसलिए प्रशिक्षण के दौरान उनके बीच स्विच करना (जैसे प्रीट्रेनिंग के साथ Muon और फाइनट्यूनिंग के साथ Muonp) किसी स्टेट रीसेट की आवश्यकता नहीं रखता, केवल बहुपद अनुमान (polynomial approximation) में बदलाव की आवश्यकता होती है।

मुख्य योगदान (Key Contributions)

1. Muonp ऑप्टिमाइज़र: बिना स्पष्ट SVD के किसी भी परिमेय $p \in (0, 1)$ के लिए $US^pV^\top$ कंप्यूट करने वाले एक सामान्यीकृत ऑप्टिमाइज़र का परिचय।
2. सैद्धांतिक सीमाएँ: प्रमाण कि भिन्नात्मक स्पेक्ट्रल शक्तियाँ निश्चित यूनिवेरिएट बहुपद पुनरावृत्तियों के माध्यम से कंप्यूट नहीं की जा सकती हैं, जिसके लिए एक नवीन द्विचरीय (bivariate) दृष्टिकोण की आवश्यकता होती है।
3. रचनात्मक एल्गोरिदम: निम्न-डिग्री विषम द्विचरीय पुनरावृत्तियों का व्युत्पन्न जो Muon की मैट्रिक्स-मल्टीप्लिकेशन-ओनली संरचना को बनाए रखते हुए भिन्नात्मक शक्तियों का अनुमान लगाते हैं।
4. अनुभवजन्य सत्यापन (Empirical Validation): गणितीय तर्क, कोडिंग और डाउनस्ट्रीम कार्यों में Llama3-1B और Llama3-3B मॉडल्स को फाइनट्यून करते समय Muonp द्वारा Muon और AdamW की तुलना में निरंतर सुधार का प्रदर्शन।
5. स्पेक्ट्रल ज्यामिति विश्लेषण: यह अंतर्दृष्टि कि इष्टतम स्पेक्ट्रल ज्यामिति शासन-निर्भर (regime-dependent) है: एकरूप अपडेट ($p=0$) प्रीट्रेनिंग के लिए फायदेमंद होते हैं, जबकि आंशिक संरक्षण ($0 < p < 1$) फाइनट्यूनिंग के लिए फायदेमंद होता है।

प्रयोगात्मक परिणाम (Experimental Results)

फाइनट्यूनिंग प्रदर्शन

Llama3-1B और Llama3-3B मॉडल्स को फाइनट्यून करते समय Muonp, Muon और AdamW की तुलना में निरंतर सुधार प्रदर्शित करता है:

• गणित और कोड रीजनिंग: GSM8K, Numina, और OpenCodeInstruct पर, Muonp, Muon की तुलना में pass@8 सटीकता में सुधार करता है और लॉस (loss) को कम करता है।
• परप्लेक्सिटी (Perplexity): Muonp, Fineweb और Pile of Law डेटासेट्स पर कम वैलिडेशन परप्लेक्सिटी प्राप्त करता है।
• डाउनस्ट्रीम टास्क: Muonp, Muon और AdamW दोनों की तुलना में रीजनिंग बेंचमार्क (PIQA, Winogrande, LAMBADA, आदि) के एक समूह में उच्च सटीकता प्रदान करता है।
• ऑप्टिमाइज़र मिसमैच: यहाँ तक कि जब प्रीट्रेनिंग के दौरान Muon (मैचिंग ऑप्टिमाइज़र) का उपयोग किया जाता है, तब भी Muonp फाइनट्यूनिंग के दौरान Muon से बेहतर प्रदर्शन करता है, जो सुझाव देता है कि यह लाभ स्पेक्ट्रल ज्यामिति के मौलिक होने के कारण है न कि केवल ऑप्टिमाइज़र निरंतरता के कारण।

प्रीट्रेनिंग बनाम फाइनट्यूनिंग द्वंद्व (Dichotomy)

पेपर प्रशिक्षण चरण के आधार पर प्रदर्शन का एक स्पष्ट द्वंद्व प्रकट करता है:

• प्रीट्रेनिंग: Muon ($p=0$), Muonp से बेहतर प्रदर्शन करता है। लेखक तर्क देते हैं कि प्रीट्रेनिंग व्यापक फीचर्स खोजने और प्रमुख दिशाओं को अपडेट पर एकाधिकार करने से रोकने के लिए "आइसोट्रोपाइजेशन" (isotropization) से लाभ उठाती है।
• फाइनट्यूनिंग: Muonp ($0 < p < 1$), Muon से बेहतर प्रदर्शन करता है। इस शासन में, मॉडल का आंतरिक आयाम (intrinsic dimension) कम होता है, और सिंगुलर वैल्यू जानकारी को बनाए रखना सबसे प्रासंगिक अवशिष्ट दिशाओं पर ध्यान केंद्रित करने में मदद करता है।
• करिकुलम प्रयोग: प्रीट्रेनिंग के अंतिम चरणों में Muon से Muonp में स्विच करना (एक करिकुलम) तत्काल लॉस रिडक्शन देता है, जो इस परिकल्पना की पुष्टि करता कि $p$ को बढ़ाने से छोटे सिंगुलर दिशाओं के लिए प्रभावी लर्निंग रेट अनुकूलित होता है।

रनटाइम (Runtime)

Muonp, Muon के तुल्य रनटाइम बनाए रखता है। उदाहरण के लिए, $p=1/3$ (क्यूबिक पुनरावृत्ति) के साथ, Muonp, Muon की तुलना में थोड़ा तेज़ है (जो कुछ कार्यान्वयनों में $p=0$ के लिए क्विंटिक पुनरावृत्ति का उपयोग करता है), H100 GPUs पर स्टेप टाइम लगभग 100–120ms है।

महत्व और दावे (Significance and Claims)

यह पेपर Muonp को Muon ऑप्टिमाइज़र के एक सिद्धांतिक सामान्यीकरण के रूप में स्थापित करता है। इसका प्राथमिक महत्व निम्नलिखित में निहित है:

1. अंतराल को पाटना: यह Muon की एकरूपता और ग्रेडिएंट डिसेंट की परिमाण संवेदनशीलता के बीच व्यवस्थित रूप से इंटरपोलेट करने का एक तरीका प्रदान करता है, जो "स्पेक्ट्रम फ्लैटनिंग" और "परिमाण प्रतिधारण" के बीच के ट्रेड-ऑफ को हल करता है।
2. शासन-निर्भर अनुकूलन: यह प्रमाण प्रदान करता है कि "सर्वश्रेष्ठ" स्पेक्ट्रल ज्यामिति सार्वभौमिक नहीं है; प्रीट्रेनिंग के लिए पूर्ण फ्लैटनिंग की आवश्यकता होती है, जबकि फाइनट्यूनिंग के लिए आंशिक फ्लैटनिंग की आवश्यकता होती है।
3. व्यावहारिक दक्षता: यह बिना SVD की लागत के इन सैद्धांतिक लाभों को प्राप्त करता है, जो बिलियन-पैरामीटर मॉडल्स तक स्केल करने के लिए आवश्यक हार्डवेयर-फ्रेंडली, मैट्रिक्स-मल्टीप्लिकेशन-ओनली संरचना को बनाए रखता है।
4. भविष्य की दिशाएं: लेखक सुझाव देते हैं कि विशिष्ट कार्यों के लिए इष्टतम $p$ का स्वचालित चयन करना और प्रीट्रेनिंग करिकुलम में Muonp को एकीकृत करना भविष्य के अनुसंधान के लिए आशाजनक मार्ग हैं।

कार्य यह निष्कर्ष निकालता है कि सिंगुलर स्पेक्ट्रम को संरक्षित करना अनुकूलन परिदृश्यों में महत्वपूर्ण लाभ ला सकता है, और Muonp इसे प्राप्त करने के लिए एक गणनात्मक रूप से कुशल, सिद्धांत-आधारित तंत्र प्रदान करता है।