🤖 machine learning

Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization

Pro-KLShampoo एक नवीन ऑप्टिमाइज़र है जो इसके प्रीकंडीशनर्स के देखे गए "स्पाइक-एंड-फ्लैट" आइजनवैल्यू स्ट्रक्चर का लाभ उठाकर KL-Shampoo को उन्नत करता है, जिससे एक निम्न-आयामी उप-स्थान (low-dimensional subspace) पर पूर्ण स्पेक्ट्रल ट्रैकिंग को शेष दिशाओं पर ऑर्थोगोनलाइजेशन के साथ संयोजित किया जाता है, जिससे कई LLM स्केल्स में वैलिडेशन लॉस, मेमोरी दक्षता और प्रशिक्षण गति में बेहतर प्रदर्शन प्राप्त होता है।

मूल लेखक: Ruotong Sun, Ermin Wei

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

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

CC BY 4.0

मूल लेखक: Ruotong Sun, Ermin Wei

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

कल्पना कीजिए कि आप एक विशाल, जटिल रोबोट (एक लार्ज लैंग्वेज मॉडल) को मानव भाषा बोलना सिखाने की कोशिश कर रहे हैं। इसे करने के लिए, आप उसे लाखों उदाहरण दिखाते हैं और उसे गलतियाँ करने देते हैं। हर बार जब वह कोई गलती करता है, तो आप उसे सही रास्ते पर लाने के लिए एक "धक्का" (ग्रेडिएंट) देते हैं।

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

इस काम के लिए दो लोकप्रिय उपकरण हैं: KL-Shampoo और Muon।

KL-Shampoo एक मास्टर मूर्तिकार की तरह है। यह हर एक धक्के के लिए एकदम सही आकार उकेरने की कोशिश करता है। यह बहुत सटीक है लेकिन पहेली के हर एक हिस्से का आकार निकालने के लिए बहुत अधिक भारी काम (कंप्यूटेशनल पावर और मेमोरी) की आवश्यकता होती है।
Muon एक जिम्नास्ट की तरह है। यह हर हिस्से को नया आकार देने की कोशिश नहीं करता; इसके बजाय, यह बस धक्कों के मोमेंटम (momentum) को सीधा करता है, जिससे वे एक साफ, ऑर्थोगोनल (लंबवत/right-angled) तरीके से आगे बढ़ते हैं। यह तेज़ और हल्का है, लेकिन यह उन सूक्ष्म विवरणों को मिस कर सकता है जिन्हें मूर्तिकार पकड़ लेता है।

बड़ी खोज: "स्पाइक-एंड-फ्लैट" (Spike-and-Flat) पैटर्न
लेखकों ने "मूर्तिकार" के काम (KL-Shampoo प्रीकंडीशनर) का बारीकी से अध्ययन किया और एक अजीब, सुंदर पैटर्न देखा। उन्होंने पाया कि "धक्के-के-आकार" (nudge-shapes) यादृच्छिक (random) नहीं हैं। वे एक स्पाइक-एंड-फ्लैट परिदृश्य की तरह दिखते हैं:

द स्पाइक (The Spike): कुछ दिशाओं में बहुत बड़े, प्रभावी मान (values) हैं (जैसे कुछ ऊंचे पहाड़)।
द फ्लैट (The Flat): बाकी की दिशाएं लगभग एक ही ऊंचाई की हैं (जैसे एक विशाल, सपाट मैदान)।

यह रोबोट के मस्तिष्क की पहली परत से लेकर अंतिम परत तक, सभी परतों में होता है। ऐसा लगता है जैसे रोबोट वास्तव में केवल कुछ विशिष्ट दिशाओं की परवाह करता है, और बाकी सब कुछ केवल "सपाट" शोर (noise) है।

समाधान: Pro-KLShampoo
लेखकों ने एक नया ऑप्टिमाइज़र बनाया जिसे Pro-KLShampoo (प्रोजेक्टेड KL-Shampoo) कहा जाता है। यह एक हाइब्रिड टूल की तरह है जो "स्पाइक-एंड-फ्लैट" खोज का उपयोग करके दोनों दुनियाओं का सर्वश्रेष्ठ लाभ उठाता है।

यह इस प्रकार काम करता है, एक सरल उपमा का उपयोग करते हुए:

"वीआईपी लाउंज" (The VIP Lounge - सबस्पेस):
एल्गोरिदम उन "स्पाइक" दिशाओं की पहचान करता है—वे कुछ महत्वपूर्ण पहाड़। यह उन्हें एक विशेष "वीआईपी लाउंज" (एक ट्रैक किया गया सबस्पेस) में रखता है। इस लाउंज के अंदर, यह उन कुछ महत्वपूर्ण दिशाओं के लिए एकदम सही आकार प्राप्त करने के लिए भारी, सटीक मूर्तिकला उपकरण (KL-Shampoo) का उपयोग करता है। यह बाकी हिस्सों पर समय बर्बाद नहीं करता है।
"खुला मैदान" (The Open Field - कॉम्प्लीमेंट):
"फ्लैट" दिशाओं (बाकी परिदृश्य) के लिए, यह हर कंकड़ को तराशने की कोशिश करना बंद कर देता है। इसके बजाय, यह ऑर्थोगोनलाइजेशन (Orthogonalization) नामक एक चतुर ट्रिक का उपयोग करता है (जो Muon टूल से ली गई है)।
- जादुई ट्रिक: लेखकों ने गणितीय रूप से सिद्ध किया है कि यदि आप इस सपाट क्षेत्र में धक्कों को केवल "सीधा" (orthogonalize) कर देते हैं, तो यह जादुği रूप से ठीक वही बीजगणितीय परिणाम (algebraic result) प्राप्त कर लेता है जैसा कि आपने वहां भारी मूर्तिकला करने से किया होता। यह ऐसा है जैसे यह महसूस करना कि एक सपाट मैदान के लिए, आपको मानचित्र की आवश्यकता नहीं है; आपको बस एक सीधी रेखा में चलना है, और आप उसी स्थान पर पहुँच जाएंगे जहाँ आप हर समन्वय (coordinate) की गणना करने के बाद पहुँचते।

यह बेहतर क्यों है?

गति (Speed): "फ्लैट" हिस्सों के लिए भारी मूर्तिकला को अनदेखा करके और उन्हें बस सीधा करके, यह एल्गोरिदम बहुत तेज़ी से चलता है। यह कंप्यूटर के "वॉलक्लॉक" (वास्तविक समय) पर बहुत समय बचाता है।
मेमोरी (Memory): इसे फ्लैट हिस्सों के लिए विशाल, जटिल आकार स्टोर करने की आवश्यकता नहीं होती है। यह केवल छोटे "वीआईपी" आकारों और बाकी के लिए एक संख्या को स्टोर करता है। इससे कंप्यूटर की मेमोरी बचती है।
प्रदर्शन (Performance): अलग-अलग आकार के रोबोट्स (GPT-2 और LLaMA) पर परीक्षणों में, Pro-KLShampoo ने मूल KL-Shampoo की तुलना में तेजी से सीखा और कम त्रुटि दर (error rate) तक पहुँचा, जबकि कम मेमोरी का उपयोग किया।

संक्षेप में
पेपर कहता है: "हमने पाया कि आधुनिक AI प्रशिक्षण के पीछे का जटिल गणित एक सरल पैटर्न रखता है: कुछ बड़े स्पाइक्स और एक फ्लैट टेल। हमने एक नया टूल बनाया है जो स्पाइक्स के साथ अत्यधिक सावधानी बरतता है और फ्लैट टेल के साथ एक सरल, तेज़ ट्रिक का उपयोग करता है। यह ट्रिक उतनी ही अच्छी तरह काम करती है जितनी कि कठिन गणित, लेकिन यह बहुत तेज़ और सस्ती है।"

परिणामस्वरूप, यह AI मॉडल्स को प्रशिक्षित करने का एक स्मार्ट, तेज़ तरीका है जो गुणवत्ता से समझौता किए बिना समय और कंप्यूटर संसाधनों की बचत करता है।

तकनीकी सारांश: Pro-KLShampoo

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

आधुनिक लार्ज लैंग्वेज मॉडल (LLM) प्री-ट्रेनिंग मुख्य रूप से उन ऑप्टिमाइज़र्स पर निर्भर करती है जो ग्रेडिएंट्स की मैट्रिक्स संरचना का लाभ उठाते हैं। दो अलग-अलग सीमाएं (frontiers) उभरी हैं:

स्पष्ट क्रोनेकर-फैक्टर्ड प्रीकंडीशनिंग (Explicit Kronecker-factored preconditioning): जिसे KL-Shampoo द्वारा दर्शाया गया है, जो वास्तविक ग्रेडिएंट कोवेरियन्स और एक क्रोनेकर-स्ट्रक्चर्ड सन्निकटन (approximation) के बीच कुलबैक-लीब्लर (KL) डाइवर्जेंस को कम करके प्रीकंडीशनर का अनुमान लगाता है।
ग्रेडिएंट मोमेंटम का ऑर्थोगोनलाइजेशन (Orthogonalization of gradient momentum): जिसका उदाहरण Muon है, जो स्पेक्ट्रल नॉर्म के तहत स्टैपेस्ट डिसेंट (steepest descent) को साकार करने के लिए मोमेंटम को ऑर्थोगोनलाइज करता है।

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

कार्यप्रणाली: Pro-KLShampoo

लेखक Pro-KLShampoo (ऑर्थोगोनलाइजेशन द्वारा व्हाइटनिंग रिकवर्ड के साथ प्रोजेक्टेड KL-Shampoo) का प्रस्ताव करते हैं। मुख्य नवाचार KL-Shampoo के प्रीकंडीशनर्स के संरचनात्मक अवलोकन से उत्पन्न होता है: उनके आइगेनवैल्यू स्पेक्ट्रा एक "स्पाइक-एंड-फ्लैट" (spike-and-flat) आकार प्रदर्शित करते हैं। इसमें कुछ प्रभावी आइगेनवैल्यू (स्पाइक) और उसके बाद एक लगभग समान टेल (फ्लैट) शामिल है। यह संरचना विभिन्न परतों और प्रशिक्षण चरणों में अनुभवजन्य रूप से देखी गई है और रैंक- $\rho$ सिग्नल-प्लस-नॉइज़ ग्रेडिएंट मॉडल के तहत सैद्धांतिक रूप से सटीक है।

1. संरचनात्मक प्रतिबंध (Structural Restriction)

Pro-KLShampoo, KL-Shampoo के एक क्रोनेकर फैक्टर (विशेष रूप से आयताकार परतों के लिए दायां-पक्ष वाला फैक्टर $R$ , जहाँ $m \le n$ ) को स्पाइक-एंड-फ्लैट संरचना के अनुरूप एक पैरामीट्रिक परिवार तक सीमित कर देता है:

प्रोजेक्टेड सबस्पेस ( $r$ -आयामी): प्रीकंडीशनर $U \in \text{St}(n, r)$ द्वारा स्पैन किए गए ट्रैक किए गए सबस्पेस पर पूर्ण स्पेक्ट्रल संरचना बनाए रखता है।
कॉम्प्लीमेंट सबस्पेस ( $n-r$ आयाम): शेष दिशाओं को एक साझा आइगेनवैल्यू $\mu_\perp$ में समाहित (collapse) कर दिया जाता है।

यह प्रतिबंध दाएं-पक्ष वाले फैक्टर के स्टोरेज को $O(n^2)$ से घटाकर $O(nr)$ और QR डिकम्पोजिशन लागत को $O(n^3)$ से घटाकर $O(r^3)$ कर देता है।

2. ऑर्थोगोनलाइजेशन के माध्यम से रिकवरी (Recovery via Orthogonalization)

यदि टेल पूरी तरह से यूनिफॉर्म नहीं है, तो एक स्केलर टेल तक सीमित होना सूचना की हानि कर सकता है। कॉम्प्लीमेंट सबस्पेस पर पूर्ण KL-Shampoo प्रीकंडीशनर के बीजगणितीय रूप को पुनः प्राप्त करने के लिए, Pro-KLShampoo कॉम्प्लीमेंट ग्रेडिएंट घटक पर ऑर्थोगोनलाइजेशन (पोलर डिकम्पोजिशन) लागू करता है।

अपडेट नियम है:
$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$
जहाँ:

पहला पद ट्रैक किए गए सबस्पेस के भीतर पूर्ण KL-Shampoo अपडेट का प्रतिनिधित्व करता है।
दूसरा पद कॉम्प्लीमेंट पर ऑर्थोगोनलाइजेशन लागू करता है।
लेखक एक पहचान (इक्वेशन 12) सिद्ध करते हैं जो यह दर्शाती है कि यह ऑर्थोगोनलाइजेशन कॉम्प्लीमेंट पर पूर्ण KL-Shampoo के प्रति-दिशा व्हाइटनिंग (per-direction whitening) के बीजगणितीय रूप को पुनः प्राप्त करता है, जो प्रभावी रूप से Muon-शैली के ऑर्थोगोनलाइजेशन को KL एस्टीमेशन फ्रेमवर्क से जोड़ता है।

3. सबस्पेस ट्रैकिंग (Subspace Tracking)

इष्टतम सबस्पेस $U$ को मैट्रिक्स $\Phi_L = E[G^\top L^{-1} G]$ के आइजनस्पेस के रूप में पहचाना गया है। स्पाइक-एंड-फ्लैट अनुमान के तहत, $\Phi_L$ का टॉप- $r$ आइजनस्पेस इष्टतम विकल्प है। एल्गोरिदम इस सबस्पेस को $\Phi_L$ के रनिंग एस्टीमेट पर पावर इटरेशन (QR डिकम्पोजिशन के माध्यम से) के एक एकल चरण का उपयोग करके ट्रैक करता है।

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

संरचनात्मक अवलोकन और पैरामीट्रिक प्रतिबंध: शोध पत्र KL-Shampoo के प्रीकंडिशनर्स में "स्पाइक-एंड-फ्लैट" आइगेनवैल्यू संरचना की पहचान करता है (जो अनुभवजन्य रूप से मजबूत और लो-रैंक ग्रेडिएंट मॉडल के तहत सैद्धांतिक रूप से सटीक है) और इसे एक पैरामीट्रिक परिवार पर KL ऑब्जेक्टिव को प्रतिबंधित करके इसका लाभ उठाता है।
ऑर्थोगोनलाइजेशन के माध्यम से बीजगणितीय रिकवरी: यह प्रदर्शित करता है कि कॉम्प्लीमेंट सबस्पेस पर ऑर्थोगोनलाइजेशन लागू करने से पूर्ण KL-Shampoo प्रीकंडीशनर का बीजगणितीय रूप पुनः प्राप्त होता है। यह Muon-शैली के ऑर्थोगोनलाइजेशन और KL एस्टीमेशन के बीच एक सैद्धांतिक सेतु प्रदान करता है।
अभिसरण गारंटी (Convergence Guarantee): लेखक ऑपरेटर-नॉर्म स्मूथनेस के तहत आइडियलाइज्ड एल्गोरिदम के लिए $O(1/\sqrt{T})$ की दर के साथ नॉन-कॉन्वेक्स अभिसरण गारंटी प्रदान करते हैं।
अनुभवजन्य प्रदर्शन (Empirical Performance): GPT-2 (124M/350M) और LLaMA (134M/450M) मॉडल्स पर व्यापक मूल्यांकन दिखाता है कि Pro-KLShampoo सभी परीक्षण किए गए सबस्पेस रैंकों ( $r \in \{32, 64, 128\}$ ) में वैलिडेशन लॉस, पीक प्रति-GPU मेमोरी और मैच की गई लॉस लेवल तक पहुँचने के लिए लगने वाले वॉलक्लॉक टाइम में KL-Shampoo से लगातार बेहतर प्रदर्शन करता है।

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

लेखकों ने चार प्री-ट्रेनिंग कॉन्फ़िगरेशन पर AdamW, Muon और KL-Shampoo के विरुद्ध Pro-KLShampoo का मूल्यांकन किया।

वैलिडेशन लॉस: Pro-KLShampoo ने सभी परीक्षण किए गए रैंक और मॉडल स्केल पर सबसे कम वैलिडेशन लॉस प्राप्त किया।
मेमोरी दक्षता: पीक प्रति-GPU मेमोरी उपयोग KL-Shampoo से कम और Muon के तुल्य था।
वॉलक्लॉक टाइम: Pro-KLShable विशिष्ट वैलिडेशन लॉस स्तरों तक KL-Shampoo की तुलना में तेजी से पहुँचा।
- GPT-2 पर, स्पीडअप मुख्य रूप से कम प्रति-स्टेप समय ( $O(n^3)$ QR को $O(r^3)$ से बदलने के कारण) द्वारा संचालित था।
- LLaMA पर, स्पीडअप प्रति-स्टेप दक्षता के अलावा तेज़ अभिसरण (कम ट्रेनिंग लॉस कर्व्स) द्वारा संचालित था।
एब्लेशन स्टडीज (Ablation Studies):
- ऑर्थोगोनलाइजेशन को हटाने (Smok-Hop वेरिएंट) से प्रदर्शन में गिरावट आई, विशेष रूप से बड़े स्केल (LLaMA 450M) पर, जो प्रति-दिशा व्हाइटनिंग को पुनः प्राप्त करने के लिए ऑर्थोगोनलाइजेशन की आवश्यकता की पुष्टि करता है।
- न तो केवल सबस्पेस-ओनली और न ही केवल कॉम्प्लीमेंट-ओनली वेरिएंट, पूर्ण Pro-KLShampoo के प्रदर्शन से मेल खा सके, जो यह दर्शाता है कि दोनों घटक आवश्यक हैं।

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

शोध पत्र का दावा है कि Pro-KLShampoo सफलतापूर्वक दो पहले से अलग-थलग पड़े ऑप्टिमाइजेशन फ्रंटियर्स को एकीकृत करता है। ग्रेडिएंट कोवेरियन्स के अंतर्निहित "स्पाइक-एंड-फ्लैट" संरचना का लाभ उठाकर, यह KL-Shampoo की सांख्यिकीय दक्षता प्राप्त करता है जबकि कम्प्यूटेशनल दक्षता बनाए रखने के लिए ऑर्थोगोनलाइजेशन का उपयोग करता है।

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

शोध पत्र अपनी सीमाओं के संबंध में विनम्र रहता है, यह नोट करते हुए कि रैंक $r$ को अनुभवजन्य रूप से सेट किया गया है, और मिक्सिंग वेट $\alpha_{kl}$ के लिए प्रति-लेयर एडेप्टिव कैलिब्रेशन तथा न्यूटन-शुलज़ इटरेशन और नेस्टरोव मोमेंटम वाले व्यावहारिक संस्करण के अभिसरण विश्लेषण के लिए भविष्य के कार्य की आवश्यकता है।