Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

तकनीकी सारांश: "डिनोइज़ फर्स्ट, ऑर्थोगोनलाइज़ लेटर: स्पेक्ट्रल फ़िल्टरिंग के माध्यम से म्यूऑन में मोमेंटम को समझना"

समस्या विवरण (Problem Statement)
म्यूऑन (Muon) ऑप्टिमाइज़र ने हाल ही में न्यूटन-शुल्ज़ इटरेशन (Newton–Schulz iteration) का उपयोग करके मोमेंटम मैट्रिक्स को लगभग ऑर्थोगोनलाइज (orthogonalize) करके, बड़े लैंग्वेज मॉडल्स (LLMs) के प्रशिक्षण में मजबूत अनुभवजन्य प्रदर्शन प्रदर्शित किया है। हालाँकि, म्यूऑन के भीतर मोमेंटम की सैद्धांतिक भूमिका अभी भी अस्पष्ट है। मौजूदा विश्लेषण या तो स्पेक्ट्रल अपडेट्स का अलग से अध्ययन करने के लिए मोमेंटम को हटा देते हैं, या इसे बिना इसके व्यावहारिक लाभों को समझाए बरकरार रखते हैं। इसके अलावा, म्यूऑन पर मौजूदा सैद्धांतिक साहित्य अक्सर यह सुझाव देता है कि भारी मोमेंटम अभिसरण दरों (convergence rates) को खराब कर सकता है या कोई लाभ नहीं देता है, जो व्यावहारिक अवलोकनों के विपरीत है जहाँ मोमेंटम आवश्यक है। इसके अतिरिक्त, ऑपरेशन्स का विशिष्ट क्रम—मोमेंटम को ऑर्थोगोनलाइजेशन से पहले लागू करना (Pre-polar) बनाम ऑर्थोगोनलाइजेशन को मोमेंटम से पहले करना (Post-polar)—एक कठोर स्पष्टीकरण की कमी रखता है, जबकि व्यावहारिक रूप से Pre-polar, Post-polar की तुलना में काफी बेहतर प्रदर्शन करता है।

कार्यप्रणाली (Methodology)
लेखक म्यूऑन का विश्लेषण करने के लिए एक संरचित सिग्नल-प्लस-परटर्बेशन ग्रेडिएंट मॉडल (signal-plus-perturbation gradient model) प्रस्तावित करते हैं। वे यह मानते हैं कि इटरेशन $t$ पर ग्रेडिएंट $G_t$ एक सुसंगत, लो-रैंक सिग्नल घटक $G^{\text{sig}}_t$ और एक बाउंडेड-वेरियंस, मीन-जीरो ऑर्थोगोनल (BVMZOS) परटर्बेशन $\Xi_t$ में विभाजित होता है।

मुख्य सैद्धांतिक योगदान मोमेंटम बफर $M_t = \beta M_{t-1} + (1-\beta)G_t$ को एक स्पेक्ट्रल फिल्टर (spectral filter) के रूप में विश्लेषण करना है। लेखक वेट अपडेट $\Delta W_t$ के लिए तीन पाइपलाइन वेरिएंट्स की जांच करते हैं:

1. Pre-polar (मानक म्यूऑन): $\Delta W_t \propto O(M_t)$, जहाँ मोमेंटम पहले लागू किया जाता है, उसके बाद पोलर फैक्टर $O(\cdot)$ लगाया जाता है।
2. Post-polar (Orthogonal-SGDM): $\Delta W_t \propto \tilde{M}_t$, जहाँ पोलर फैक्टर पहले कच्चे ग्रेडिएंट पर लागू किया जाता है, फिर मोमेंटम संचित किया जाता है।
3. Polar-only: $\Delta W_t \propto O(G_t)$, बिना मोमेंटम के एक बेसलाइन।

विश्लेषण प्रभावी विंडो साइज (effective window size) $T = 1/(1-\beta)$ पर निर्भर करता है और रैंडम मैट्रिक्स थ्योरी के उपकरणों का उपयोग करता है, जिसमें सिंगुलर सबस्पेस परटर्बेशन के लिए वेडिन का $\sin \Theta$ प्रमेय शामिल है, ताकि ग्रेडिएंट मैट्रिक्स के सिंगुलर मानों और सबस्पेसों को मापने के लिए।

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

1. मोमेंटम एक स्पेक्ट्रल फिल्टर के रूप में (Theorem 1 & Corollary 1):
   लेखक सिद्ध करते हैं कि उनके ग्रेडिएंट मॉडल के तहत, मोमेंटम एक लो-पास फिल्टर के रूप में कार्य करता है जो परटर्बेशन (शोर/noise) को दबाता है जबकि सुसंगत सिग्नल को सुरक्षित रखता है। विशेष रूप से, जैसे-जैसे प्रभावी विंडो साइज $T$ बढ़ता है, टेल सिंगुलर वैल्यूज (जो शोर द्वारा संचालित होती हैं) $O(T^{-1/4})$ की दर से घटती हैं, जबकि लीडिंग सिंगुलर वैल्यूज (सिग्नल) शून्य से दूर बनी रहती हैं। यह सिग्नल और शोर घटकों के बीच एक विस्तृत स्पेक्ट्रल गैप (spectral gap) बनाता है। फलस्वरूप, मोमेंटम बफर $M_t$ के टॉप-$r$ सिंगुलर सबस्पेस, कच्चे ग्रेडिएंट के सबस्पेस की तुलना में वास्तविक सिग्नल सबस्पेस के साथ अधिक विश्वसनीय रूप से संरेखित होते हैं।

2. नॉन-कम्यूटेटिविटी और पाइपलाइन ऑर्डरिंग (Theorem 2):
   यह शोध पत्र स्थापित करता है कि मोमेंटम और ऑर्थोगोनलाइजेशन कम्यूट (commute) नहीं होते हैं।

• Pre-polar का प्रभुत्व: जब मोमेंटम पहले लागू किया जाता है, तो यह ग्रेडिएंट स्ट्रीम को 'डिनोइज़' (denoise) करता है, जिससे बाद में पोलर फैक्टर एक स्थिर सिंगुलर सबस्पेस संरचना वाले मैट्रिक्स पर कार्य कर पाता है। लेखक सिद्ध करते हैं कि Pre-polar, Post-polar और Polar-only बेसलाइन्स दोनों की तुलना में उच्च अपेक्षित सिग्नल अलाइनमेंट प्राप्त करता है।
• विकल्पों की विफलता: Post-polar के लिए, शोर वाले ग्रेडिएंट को पहले ऑर्थोगोनलाइज करना सिग्नल और शोर की दिशाओं के बीच एम्प्लीट्यूड गैप को "मिटा" देता है। पोलर फैक्टर सभी सिंगुलर वैल्यूज को एक (ones) से बदल देता है, जिससे बाद के मोमेंटम स्टेप्स में सिग्नल को शोर से अलग करना असंभव हो जाता है। सैद्धांतिक रूप से, Post-polar और Polar-only पाइपलाइन प्रभावी ढंग से सिग्नल दिशा को रिकवर करने में विफल रहते हैं, विशेष रूप से लो सिग्नल-टू-नॉइज़ रेशियो (SNR) वाले परिदृश्यों में।

3. लो SNR में मात्रात्मक पृथक्करण (Theorem 3):
   एक रैंक-1 स्पाइक्ड गौसियन मॉडल में, लेखक दिखाते हैं कि जैसे-जैसे मैट्रिक्स डायमेंशन $m \to \infty$ होता है, Polar-only पाइपलाइन का सिग्नल अलाइनमेंट लुप्त हो जाता है। इसके विपरीत, Pre-polar पाइपलाइन तब तक लक्षित सिग्नल को रिकवर करना जारी रखती है जब तक कि मोमेंटम विंडो $T$ पर्याप्त बड़ी न हो।

परिणाम (Results)
इन सैद्धांतिक निष्कर्षों को व्यापक प्रयोगों के माध्यम से मान्य किया गया है:

• स्पेक्ट्रल गैप विज़ुअलाइज़ेशन: NanoGPT पर स्टेशनरी प्रोब्स का उपयोग करके, लेखक प्रदर्शित करते हैं कि मोमेंटम कोएफिशिएंट $\beta$ को बढ़ाने से मोमेंटम बफर के लीडिंग और टेल सिंगुलर वैल्यूज के बीच का अंतर बढ़ जाता है, जो Theorem 1 के अनुरूप है।
• सबस्पेस विश्वसनीयता: सबस्पेस अलाइनमेंट एरर के अनुभवजन्य माप दिखाते हैं कि मोमेंटम बफर के सिंगुलर सबस्पेस और वास्तविक सिग्नल सबस्पेस के बीच विचलन $T$ बढ़ने के साथ कम होता है, जो Corollary 1 को मान्य करता है।
• पाइपलाइन तुलना: सिंथेटिक मॉडल्स, CIFAR-10, और NanoGPT/LLaMA 350M ट्रेनिंग के माध्यम से, Pre-polar पाइपलाइन लगातार सिग्नल अलाइनमेंट और वैलिडेशन लॉस के मामले में Post-polar और Polar-only से बेहतर प्रदर्शन करती है। विशेष रूप से, Post-polar का प्रदर्शन $\beta$ के बावजूद स्थिर रहता है, जो इस सैद्धांतिक भविष्यवाणी की पुष्टि करता है कि औसत निकालने से पहले ऑर्थोगोनलाइज करने से शोर के कारण खोई हुई सिग्नल दिशाओं को रिकवर नहीं किया जा सकता।
• एंड-टू-एंड ट्रेनिंग: LLM प्रीट्रेनिंग कार्यों (NanoGPT और LLaMA 350M) में, Pre-polar पाइपलाइन (मानक म्यूऑन) अन्य वेरिएंट्स की तुलना में कम वैलिडेशन लॉस प्राप्त करती है, जो इस दावे का समर्थन करती है कि मोमेंटम सिग्नल रिकवरी के लिए एक आवश्यक घटक है, न कि केवल एक एक्सेलेरेशन ट्रिक।

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