A Proof of Learning Rate Transfer under $μ$P

यहाँ सुफियान हायौ के शोध पत्र "A Proof of Learning Rate Transfer under µP" का विस्तृत तकनीकी सारांश दिया गया है।

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

डीप लर्निंग में, न्यूरल नेटवर्क को स्केल करने (चौड़ाई $n$ बढ़ाने) के लिए अक्सर हाइपरपैरामीटर्स (HPs), विशेष रूप से लर्निंग रेट ($\eta$), को फिर से ट्यून करने की आवश्यकता होती है ताकि प्रशिक्षण स्थिरता और प्रदर्शन बना रहे।

• घटना (The Phenomenon): अनुभवजन्य अवलोकन (Empirical observations) सुझाव देते हैं कि मैक्सिमल अपडेट पैरामीट्रेशन (µP) के तहत, नेटवर्क की चौड़ाई बढ़ने के साथ इष्टतम लर्निंग रेट एक स्थिरांक (constant) की ओर अभिसरित (converge) होता है। यह अभ्यासकर्ताओं को छोटे मॉडलों पर हाइपरपैरामीटर्स ट्यून करने और उन्हें बिना पुन: ट्यूनिंग के बड़े मॉडलों पर स्थानांतरित करने (Zero-shot HP transfer) की अनुमति देता है।
• अंतराल (The Gap): जबकि यांग और हू (2021) ने यह अनुमान लगाया था कि यह घटना µP द्वारा अनंत-चौड़ाई सीमा (infinite-width limit) में "मैक्सिमल फीचर लर्निंग" का परिणाम है, लेकिन इसका कोई कठोर गणितीय प्रमाण नहीं था। इसके विपरीत, स्टैंडर्ड पैरामीट्रेशन (SP) और न्यूरल टेंगेंट पैरामीट्रेशन (NTP) के तहत, इष्टतम लर्निंग रेट को चौड़ाई बढ़ने के साथ बदलते हुए (आमतौर पर शून्य की ओर घटते या विचलित होते) देखा गया है।
• उद्देश्य: इस शोध पत्र का लक्ष्य गहरे रैखिक मल्टी-लेयर पर्सेप्ट्रॉन (MLPs) के लिए µP के तहत लर्निंग रेट ट्रांसफर का पहला कठोर प्रमाण प्रदान करना और यह प्रदर्शित करना है कि यह अन्य पैरामीट्राइजेशन के तहत क्यों विफल हो जाता है।

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

लेखक अनंत-चौड़ाई सीमा ($n \to \infty$) में ग्रेडिएंट डिसेंट (GD) के साथ प्रशिक्षित डीप लीनियर नेटवर्क्स के सैद्धांतिक विश्लेषण का उपयोग करते हैं।

• मॉडल सेटअप:

  • एक $L$ परतों वाला लीनियर MLP: $f(x) = V^\top W_L \dots W_1 W_0 x$।
  • प्रशिक्षण द्विघाती हानि (quadratic loss) को GD का उपयोग करके कम करता है।
  • न्यूरल पैरामीट्रेशन: यह पेपर औपचारिक रूप रूप से बताता है कि कैसे इनिशियलाइजेशन वेरिएंस और लर्निंग रेट्स, एक्सपोनेंट्स $\alpha_\ell$ और $c$ के माध्यम से चौड़ाई $n$ के साथ स्केल होते हैं।
    • µP: $W_\ell \sim \mathcal{N}(0, n^{-1})$, $V \sim \mathcal{N}(0, n^{-2})$, $\eta \sim \mathcal{O}(1)$।
    • SP: $W_\ell \sim \mathcal{N}(0, n^{-1})$, $V \sim \mathcal{N}(0, n^{-1})$, $\eta \sim \mathcal{O}(1)$।

• मुख्य विश्लेषणात्मक दृष्टिकोण:

  • पॉलीनोमियल एक्सपेंशन: मुख्य अंतर्दृष्टि यह है कि एक लीनियर नेटवर्क के लिए, $t$ प्रशिक्षण चरणों के बाद मॉडल आउटपुट $f^{(t)}(x)$ को लर्निंग रेट $\eta$ के एक पॉलीनोमियल फंक्शन के रूप में व्यक्त किया जा सकता है।
  • एसिम्प्टोटिक विश्लेषण (Asymptotic Analysis): इस पॉलीनोमियल के गुणांक (coefficients) रैंडम वेरिएबल्स हैं जो इनिशियलाइजेशन पर निर्भर करते हैं। लेखक स्ट्रॉन्ग लॉ ऑफ लार्ज नंबर्स (SLLN) और रैंडम मैट्रिक्स प्रोडक्ट्स (टेन्सर प्रोग्राम्स) के गुणों का उपयोग करके $n \to \infty$ होने पर इन गुणांकों के अभिसरण का विश्लेषण करते हैं।
  • लॉस लैंडस्केप: चूंकि लॉस आउटपुट एरर का वर्ग है, इसलिए लॉस फंक्शन $L_n(\eta)$ $\eta$ में एक पॉलीनोमियल बन जाता है। इष्टतम लर्निंग रेट इस पॉलीनोमियल का मिनिमाइज़र (minimizer) है।
  • मिनिमा का अभिसरण: प्रमाण इस बात पर निर्भर करता है कि रैंडम पॉलीनोमियल्स का अनुक्रम (sequence) एक नियतात्मक (deterministic) लिमिटिंग पॉलीनोमियल में लगभग निश्चित रूप से (almost surely) अभिसरित होता है। रैंडम पॉलीनोमियल का मिनिमाइज़र फिर लिमिटिंग पॉलीनोमियल के मिनिमाइज़र की ओर अभिसरित होता है।

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

A. µP के तहत LR Transfer का सैद्धांतिक प्रमाण

यह पेपर सिद्ध करता है कि µP के तहत, इष्टतम लर्निंग रेट $\eta_n^*$ एक गैर-शून्य नियतात्मक स्थिरांक $\eta_\infty^*$ की ओर अभिसरित होता है जब $n \to \infty$।

• चरण 1 ($t=1$): आउटपुट $f^{(1)}(x)$ को $f^{(0)}(x) + \sum \phi_\ell \eta^\ell$ के रूप में विस्तारित किया गया है।
  • µP के तहत, लीनियर टर्म का गुणांक ($\phi_1$) एक गैर-शून्य स्थिरांक की ओर अभिसरित होता है।
  • उच्च-क्रम के पदों के गुणांक ($\phi_{\ell}, \ell \ge 2$) $O(n^{-(\ell-1)/2})$ के रूप में लुप्त हो जाते हैं।
  • फलस्वरूप, लॉस फंक्शन $\eta$ के संदर्भ में अभिसारी रूप से क्वाड्रेटिक (quadratic) हो जाता है, जो एक अद्वितीय, गैर-शून्य मिनिमाइज़र सुनिश्चित करता है।
• सामान्य चरण ($t \ge 2$): लेखक इसे किसी भी प्रशिक्षण चरण तक विस्तारित करते हैं। हालांकि पॉलीनोमियल की डिग्री बढ़ जाती है (जो $\eta$ में नॉन-लीनियर हो जाती है), फिर भी लिमिटिंग पॉलीनोमियल में एक अद्वितीय मिनिमाइज़र होता है, बशर्ते कि लिमिटिंग लॉस का एक अद्वितीय मिनिमाइज़र हो।
• अभिसरण दर (Convergence Rate): इष्टतम लर्निंग रेट $O(n^{-1/2})$ की दर से अभिसरित होता है।

B. SP और NTP के तहत विफलता का प्रमाण

यह पेपर सिद्ध करता है कि स्टैंडर्ड पैरामीट्रेशन (SP) और न्यूरल टेंगेंट पैरामीट्रेशन (NTP) के तहत LR ट्रांसफर विफल हो जाता है।

• SP के तहत: आउटपुट लेयर $V$ का इनिशियलाइजेशन वेरिएंस $n^{-1}$ है (µP में $n^{-2}$ के बजाय)। इसके कारण सभी पॉलीनोमियल गुणांक $\sqrt{n}$ के साथ स्केल होते हैं। लॉस को कम करने के लिए, इष्टतम लर्निंग रेट को $1/\sqrt{n}$ (या उससे तेज़) स्केल होना चाहिए, जिससे $n \to \infty$ होने पर $\eta_n^* \to 0$ हो जाता है।
• NTP के तहत: स्केलिंग एक "लेज़ी ट्रेनिंग" (lazy training) शासन की ओर ले जाती है जहाँ फीचर्स लर्न नहीं कर पाते। इष्टरल लर्निंग रेट कर्नेल रिजीम डायनेमिक्स की भरपाई करने के लिए विचलित या महत्वपूर्ण रूप से शिफ्ट होता है।

C. अनुभवजन्य सत्यापन (Empirical Validation)

सिंथेटिक लीनियर और नॉन-लीनियर (ReLU) डेटासेट्स पर व्यापक प्रयोग सिद्धांत की पुष्टि करते हैं:

• µP: चौड़ाई $n \in [128, 8192]$ के लिए इष्टतम लर्निंग रेट एक स्थिर मान के आसपास क्लस्टर होते हैं।
• SP: चौड़ाई बढ़ने के साथ इष्टतम लर्निंग रेट शून्य की ओर बढ़ते हैं।
• मजबूती (Robustness): यह घटना विभिन्न डेप्थ्स ($L=3, 9, 27$), ऑप्टिमाइज़र्स (Adam), और ट्रेनिंग स्टेप्स (अभिसरण के करीब) के लिए बनी रहती है।

4. मुख्य परिणाम (Key Results)

1. थ्योरम 1 (Step $t=1$ पर LR Transfer): µP के लिए, $\eta_n^{(1)} \xrightarrow{P} \eta_\infty^{(1)} = \frac{m}{L} \frac{y^\top K y}{\|Ky\|^2}$, जहाँ $K$ नॉर्मलाइज्ड इनपुट ग्राम मैट्रिक्स है। अभिसरण दर $O(n^{-1/2})$ है।
2. थ्योरम 2 (SP के तहत No LR Transfer): SP के तहत, $\eta_n^{(1)} \xrightarrow{P} 0$ होता है जब $n \to \infty$।
3. थ्योरम 3 (Step $t$ पर LR Transfer): सौम्य शर्तों (लिमिटिंग लॉस के अद्वितीय मिनिमाइज़र) के तहत, µP के लिए किसी भी ट्रेनिंग स्टेप $t$ पर LR ट्रांसफर होता है। इष्टतम लर्निंग रेट लगभग निश्चित रूप से एक गैर-शून्य स्थिरांक की ओर अभिसरित होता है।
4. अनुभवजन्य अभिसरण: सिमुलेशन दिखाते हैं कि व्यावहारिक चौड़ाई (जैसे, $n=1024$) के लिए, इष्टतम LR और लिमिट के बीच का अंतर नगण्य है, जो µP की हाइपरपैरामीटर ट्रांसफर उपयोगिता को प्रमाणित करता है।

5. महत्व और प्रभाव (Significance and Impact)

• सैद्धांतिक आधार: यह कार्य µP के "जीरो-शॉट हाइपरपैरामीटर ट्रांसफर" गुण के लिए पहला कठोर गणितीय औचित्य प्रदान करता है, जो केवल अनुमानित तर्कों से आगे जाता है।
• व्यावहारिक उपयोगिता: यह छोटे मॉडलों पर लर्निंग रेट ट्यून करने और उन्हें विशाल मॉडलों (जैसे, LLMs) पर लागू करने की रणनीति को मान्य करता है, जिससे हाइपरपैरामीटर सर्च की कम्प्यूटेशनल लागत काफी कम हो जाती है।
• पैरामीट्राइजेशन का स्पष्टीकरण: यह गणितीय रूप से µP को SP और NTP से अलग करता है, और यह समझाता है कि क्यों µP फीचर लर्निंग और स्थिर स्केलिंग के लिए बेहतर है। यह रेखांकित करता है कि आउटपुट लेयर का विशिष्ट इनिशियलाइजेशन ($V \sim n^{-2}$) ही स्थिर फीचर अपडेट और स्थिर इष्टतम लर्निंग रेट को सक्षम करने वाला महत्वपूर्ण कारक है।
• सीमाएँ और भविष्य का कार्य: वर्तमान प्रमाण रैखिक नेटवर्क और ग्रेडिएंट डिसेंट तक सीमित है। लेखक नोट करते हैं कि नॉन-लीनियर नेटवर्क और जटिल ऑप्टिमाइज़र्स (जैसे मोमेंटम के साथ Adam) तक विस्तार करने के लिए नई प्रूफ मशीनरी की आवश्यकता होगी, हालांकि अनुभवजन्य परिणाम बताते हैं कि यह घटना बनी रहती है।

संक्षेप में, यह पेपर स्थापित करता है कि µP एकमात्र ऐसा पैरामीट्रेशन (मानक विकल्पों में से) है जो यह सुनिश्चित करता है कि नेटवर्क की चौड़ाई अनंत तक बढ़ने पर इष्टतम लर्निंग रेट स्थिर रहे, जो आधुनिक बड़े पैमाने के AI प्रशिक्षण में µP की दक्षता के लिए एक ठोस सैद्धांतिक आधार प्रदान करता है।