LAMB ऑप्टिमाइज़र (लेयर-वाइज़ एडेप्टिव मोमेंट्स फॉर बैच ट्रेनिंग) एक स्टोकेस्टिक ऑप्टिमाइज़ेशन एल्गोरिदम है जो गहरे तंत्रिका नेटवर्क को प्रशिक्षित करने के लिए उपयोग किया जाता है, और यह एडम ऑप्टिमाइज़र को प्रति-लेयर सामान्यीकरण चरण के साथ विस्तारित करता है। 2019 में गूगल के शोधकर्ताओं (विशेष रूप से यांग यू, जिंग ली, जोनाथन हसू, और अन्य) द्वारा पेश किया गया, LAMB को बहुत बड़े मिनी-बैच आकारों (जैसे, 32,768 या अधिक) के कुशल उपयोग को सक्षम करने के लिए डिज़ाइन किया गया था, बिना मॉडल सटीकता को कम किए या व्यापक हाइपरपैरामीटर ट्यूनिंग की आवश्यकता के। यह प्रत्येक परत के वजन मानक और उसके अपडेट मानक के अनुपात के आधार पर प्रत्येक परत के लिए अपडेट परिमाण को स्केल करके इसे प्राप्त करता है, प्रभावी रूप से सीखने की दर को प्रत्येक परत के पैरामीटर के पैमाने से अलग करता है।
LAMB विशेष रूप से Transformer (architecture)-आधारित मॉडल को प्रशिक्षित करने में प्रभावशाली रहा है, जिसमें शुरुआती बड़े भाषा मॉडल (LLM) और ResNet जैसे विज़न आर्किटेक्चर शामिल हैं। वितरित प्रशिक्षण ढांचे में इसका अपनाना, जैसे TensorFlow (tf.keras.optimizers.LAMB के माध्यम से) और PyTorch (NVIDIA, Hugging Face, और FairScale जैसी लाइब्रेरीज़ में LAMB कार्यान्वयन के माध्यम से), ने इसे प्रशिक्षण रन को स्केल करने के लिए एक मानक उपकरण बना दिया। बड़े बैचों की अनुमति देकर, LAMB अत्याधुनिक मॉडलों को प्रशिक्षित करने के लिए आवश्यक दीवार-घड़ी के समय को कम करता है, जो OpenAI, Anthropic, और Google DeepMind जैसे संगठनों के लिए महत्वपूर्ण है जो विशाल कंप्यूट क्लस्टर पर निर्भर हैं।
प्रेरणा और पृष्ठभूमि
बड़े मिनी-बैच के साथ गहरे तंत्रिका नेटवर्क को प्रशिक्षित करने से प्रति-युग कम्प्यूटेशनल चरणों की संख्या कम हो जाती है, लेकिन बैच आकार का भोला-भाला स्केलिंग अक्सर खराब सामान्यीकरण और अस्थिर अभिसरण की ओर ले जाता है। इसे "बड़े-बैच प्रशिक्षण समस्या" के रूप में जाना जाता है। स्टोकेस्टिक ग्रेडिएंट डिसेंट (SGD) जैसे मोमेंटम या एडम के साथ मानक ऑप्टिमाइज़र को बैच आकार बढ़ने पर सीखने की दर के सावधानीपूर्वक समायोजन की आवश्यकता होती है, और फिर भी, सटीकता अक्सर कम हो जाती है। LAMB ऑप्टिमाइज़र को बैच आकार परिवर्तनों के प्रति ऑप्टिमाइज़र को अधिक मजबूत बनाकर इस समस्या का समाधान करने के लिए विकसित किया गया था।
LAMB के पीछे मुख्य अंतर्दृष्टि यह है कि एक गहरे नेटवर्क में विभिन्न परतें ग्रेडिएंट और वजन मानक के बहुत भिन्न पैमाने प्रदर्शित करती हैं। उदाहरण के लिए, शुरुआती कन्वोल्यूशनल परतों में छोटे वजन होते हैं, जबकि बाद की पूरी तरह से जुड़ी परतों में बड़े होते हैं। एडम में एक एकल वैश्विक सीखने की दर कुछ परतों के लिए बहुत बड़े अपडेट (विचलन पैदा करने) या दूसरों के लिए बहुत छोटे अपडेट (अभिसरण को धीमा करने) का कारण बन सकती है। LAMB एक लेयर-वाइज़ एडेप्टिव रेट पेश करता है जो परत के वजन और ग्रेडिएंट मानक के अनुपात के आधार पर अपडेट को सामान्य करता है, यह सुनिश्चित करता है कि प्रत्येक परत अपने परिमाण के सापेक्ष स्थिर गति से चलती है।
एल्गोरिदम विवरण
LAMB को अतिरिक्त सामान्यीकरण चरण के साथ एडम के एक प्रकार के रूप में देखा जा सकता है। मान लें \(\theta_t\) पुनरावृत्ति \(t\) पर पैरामीटर हैं, और \(g_t\) \(\theta_t\) के संबंध में हानि का ग्रेडिएंट है। LAMB ग्रेडिएंट के पहले और दूसरे क्षणों (\(m_t\) और \(v_t\)) को बनाए रखता है, एडम के समान, घातीय क्षय दरों \(\beta_1\) और \(\beta_2\) (आमतौर पर 0.9 और 0.999) के साथ। पूर्वाग्रह सुधार के बाद, यह एक एडम अपडेट \(\frac{m_t}{\sqrt{v_t} + \epsilon}\) की गणना करता है।
महत्वपूर्ण अंतर विश्वास अनुपात \(\phi\) है: प्रत्येक परत \(i\) के लिए, \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), जहां \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) उस परत के लिए एडम अपडेट है (बिना सीखने की दर के), और \(||\cdot||\) L2 मानक को दर्शाता है। परत \(i\) के लिए अंतिम अपडेट है \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), जहां \(\eta\) वैश्विक सीखने की दर है। यह विश्वास अनुपात अपडेट को परत के वजन मानक के अनुपात में स्केल करता है, इसलिए छोटी परतों को छोटे (लेकिन नगण्य नहीं) अपडेट मिलते हैं, और बड़ी परतों को बड़े, फिर भी स्थिर, अपडेट मिलते हैं।
व्यवहार में, शून्य से विभाजन से बचने के लिए हर में एक छोटा स्थिरांक (जैसे, 1e-6) जोड़ा जाता है। एल्गोरिदम में वैकल्पिक वजन क्षय (L2 नियमितीकरण) भी शामिल है जो अपडेट में फ़्यूज़ किया जाता है, एडमडब्ल्यू में उपयोग किए गए डिकपल्ड वजन क्षय दृष्टिकोण का पालन करते हुए। जब सभी परतों के लिए विश्वास अनुपात 1 पर सेट किया जाता है, तो LAMB मानक एडम (पूर्वाग्रह सुधार के साथ) में कम हो जाता है।
हाइपरपैरामीटर और ट्यूनिंग
LAMB एडम से अधिकांश हाइपरपैरामीटर विरासत में लेता है: \(\beta_1\) (मोमेंटम), \(\beta_2\) (वेरिएंस क्षय), \(\epsilon\) (संख्यात्मक स्थिरता), और वजन क्षय दर। प्राथमिक नया हाइपरपैरामीटर वैश्विक सीखने की दर \(\eta\) है, जो अक्सर बड़े-बैच प्रशिक्षण के लिए 0.01-0.1 की सीमा में सेट किया जाता है, जो एडम के लिए विशिष्ट (जैसे, 1e-3) से काफी अधिक है। लेखकों ने पाया कि बहुत बड़े बैचों (जैसे, BERT के लिए 32,768) के लिए, पहले 10% चरणों में रैखिक वार्मअप के साथ 0.01 की सीखने की दर अच्छी तरह से काम करती है, और उन्होंने शेष चरणों पर सीखने की दर अनुसूची के रूप में कोसाइन क्षय का उपयोग करने की भी सिफारिश की।
इसके अतिरिक्त, \(\beta_2\) का चुनाव स्थिरता को प्रभावित कर सकता है; स्पार्स ग्रेडिएंट वाले मॉडल के लिए, उच्च \(\beta_2\) (जैसे, 0.99) का उपयोग किया जा सकता है। लेखकों ने यह भी सुझाव दिया कि बैच आकार को सीखने की दर के अनुपात में स्केल किया जा सकता है (रैखिक स्केलिंग नियम), एक दिशानिर्देश जो LAMB के साथ अच्छी तरह से काम करता है। उदाहरण के लिए, यदि बैच आकार दोगुना है, तो सीखने की दर भी सटीकता के नुकसान के बिना दोगुनी की जा सकती है।
प्रदर्शन और बेंचमार्क
मूल पेपर में, LAMB का मूल्यांकन दो मुख्य कार्यों पर किया गया था: इमेजनेट पर ResNet-50 का प्रशिक्षण (छवि वर्गीकरण) और मास्क किए गए भाषा मॉडलिंग के लिए BERT (एक ट्रांसफार्मर-आधारित भाषा मॉडल)। LAMB का उपयोग करके, लेखकों ने 32,768 के बैच आकार के साथ केवल 2,048 पुनरावृत्तियों में इमेजनेट शीर्ष -1 सटीकता 76.0% हासिल की, जो कम बैचों (जैसे, 256) के साथ बहुत कम युगों में प्राप्त अत्याधुनिक सटीकता से मेल खाती है। BERT के लिए, उन्होंने मॉडल को 1,024 TPU का उपयोग करके लगभग 3.5 मिनट में समान सटीकता (जैसे, SQUAD पर F1-स्कोर 1.0) के लिए प्रशिक्षित किया, जो पिछले तरीकों की तुलना में 10 गुना गति है।
इसके बाद, LAMB गूगल की आंतरिक वर्कफ़्लो में BERT-आधारित मॉडल को प्रशिक्षित करने के लिए डिफ़ॉल्ट ऑप्टिमाइज़र बन गया। पेपर ने बताया कि बैच आकार को 1,024 से 65,536 तक स्केल करते समय LAMB ने एडम और मोमेंटम के साथ SGD दोनों से बेहतर प्रदर्शन किया। लेखकों ने यह भी दिखाया कि LAMB Gradient Clipping के साथ अच्छी तरह से काम करता है (विस्फोटक ग्रेडिएंट को रोकने के लिए उपयोग किया जाता है) और मिश्रित-परिशुद्धता प्रशिक्षण के साथ संगत है, जैसा कि NVIDIA GPU और गूगल TPU जैसे आधुनिक हार्डवेयर पर उपयोग किया जाता है।
बड़े पैमाने पर प्रशिक्षण में अनुप्रयोग
LAMB का प्राथमिक अनुप्रयोग वितरित प्रशिक्षण में है जहां बैच आकार एक एकल डिवाइस की मेमोरी में फिट होने के लिए बहुत बड़ा है। ऐसे सेटअप में, ग्रेडिएंट को डेटा समानांतरता का उपयोग करके कई GPU या TPU में औसत किया जाता है। उदाहरण के लिए, OpenAI और Google DeepMind बड़े Transformer (architecture) मॉडल को प्रशिक्षित करते समय LAMB के समान ऑप्टिमाइज़र का उपयोग करते हैं जिसमें अनुक्रम लंबाई हजारों में होती है। हालांकि LAMB (और इसके उत्तराधिकारी LAMB2) जैसे नए ऑप्टिमाइज़र प्रस्तावित किए गए हैं, LAMB कई ओपन-सोर्स प्रयासों में एक विश्वसनीय विकल्प बना हुआ है, जिसमें विज़न ट्रांसफार्मर और LLM का प्रशिक्षण शामिल है जो AI21 Labs और SambaNova जैसी शोध टीमों और कंपनियों द्वारा किया जाता है।
मशीन लर्निंग के संदर्भ में Amazon Web Services पर (AWS Trainium हार्डवेयर के साथ), LAMB को दक्षता के लिए कस्टम कर्नेल में समर्थित किया जाता है। इसी तरह, Intel और AMD ने अपने त्वरक पर LAMB का बेंचमार्क किया है। चरम बैच आकार को संभालने की ऑप्टिमाइज़र की क्षमता इसे विशाल डेटासेट पर मॉडल को पूर्व-प्रशिक्षित करने के लिए मूल्यवान बनाती है, जहां एक एकल युग की लागत अधिक होती है, और युगों को कम करना सर्वोपरि है।
अन्य ऑप्टिमाइज़र से संबंध
LAMB एडेप्टिव ऑप्टिमाइज़र के परिवार का हिस्सा है जिसमें SGD प्रकार, एडम, और एडमडब्ल्यू (डिकपल्ड वजन क्षय) और LARS (लेयर-वाइज़ एडेप्टिव रेट स्केलिंग) जैसे इसके उत्तराधिकारी शामिल हैं। LARS, जिसे यू एट अल द्वारा 2017 में सीएनएन के बड़े-बैच प्रशिक्षण के लिए पेश किया गया था, समान लेयर-वाइज़ विश्वास अनुपात का उपयोग करता है लेकिन दूसरे क्षणों को बनाए नहीं रखता है; यह पहले क्षणों (मोमेंटम) और ग्रेडिएंट मानक पर निर्भर करता है। LAMB LARS के लाभों (लेयर-वाइज़ स्केलिंग) को एडम की एडेप्टिव प्रति-पैरामीटर सीखने की दरों के साथ जोड़ता है, जिससे यह स्पार्स ग्रेडिएंट वाले मॉडल (जैसे ट्रांसफार्मर) के लिए अधिक मजबूत हो जाता है।
एक और निकट से संबंधित ऑप्टिमाइज़र NVLAMB है (Nvidia से), जो वेरिएंस रिडक्शन को शामिल करता है। हालांकि, LAMB सरल और व्यापक रूप से उपयोग किया जाने वाला बना हुआ है। Sequence-to-Sequence (Seq2Seq) कार्यों के लिए Beam Search के साथ, LAMB सीधे अनुमान को प्रभावित नहीं करता है, लेकिन यह प्रशिक्षण अभिसरण में मदद करता है, जो अप्रत्यक्ष रूप से अनुक्रम डिकोडिंग में सुधार करता है।
विस्तार और प्रकार
अपनी शुरुआत के बाद से, कई प्रकार प्रस्तावित किए गए हैं। LAMB2 (गूगल से भी) ग्रेडिएंट वेरिएंस के आधार पर एक सामान्यीकरण कारक जोड़ता है, कुछ समस्याओं के लिए स्थिरता में सुधार करता है। पूर्वाग्रह सुधार के साथ LARS और अन्य संशोधन भी आम हैं। व्यवहार में, कई फ्रेमवर्क क्षणों के लिए वैकल्पिक पूर्वाग्रह सुधार के साथ LAMB को लागू करते हैं, जो पहले कुछ चरणों के दौरान फायदेमंद है। कुछ कार्यान्वयन, जैसे PyTorch के torch.optim.Lamb (torch_optimizer पैकेज में), विश्वास अनुपात पैरामीटर को समायोजित करने या कस्टम प्रति-लेयर सीखने की दर का उपयोग करने की अनुमति देते हैं।
विभिन्न स्केलिंग रणनीतियों (जैसे, 1cycle शेड्यूल) के साथ एडमडब्ल्यू जैसे नए ऑप्टिमाइज़र के उभरने के बावजूद, LAMB बड़े-बैच प्रशिक्षण के लिए एक मजबूत आधार रेखा बना हुआ है। शोध समुदाय ने सामान्यीकरण में और सुधार करने के लिए LAMB को डेटा वृद्धि और ग्रेडिएंट क्लिपिंग के साथ संयोजित करने का पता लगाया है।
व्यावहारिक विचार और सीमाएं
जबकि LAMB बड़े-बैच सेटिंग्स में उत्कृष्ट है, यह हमेशा छोटे बैच आकार (जैसे, 1,024 से नीचे) के लिए सबसे अच्छा विकल्प नहीं है। ऐसे क्षेत्रों में, मानक एडम या मोमेंटम के साथ SGD सरल और समान रूप से प्रभावी हो सकता है। LAMB प्रति परत मानकों की गणना का एक कम्प्यूटेशनल अधिभार जोड़ता है, जो आधुनिक हार्डवेयर पर नगण्य है लेकिन कई छोटी परतों वाले मॉडल (जैसे, U-Net आर्किटेक्चर) के लिए गैर-तुच्छ हो सकता है।
एक और सीमा यह है कि LAMB का विश्वास अनुपात कभी-कभी अस्थिर प्रशिक्षण का कारण बन सकता है यदि कुछ परतों में बहुत छोटे वजन मानक (शून्य के पास) हों। यह आमतौर पर हर में एक एप्सिलॉन शब्द जोड़कर और वजन क्षय का उपयोग करके कम किया जाता है, जो वजन को शून्य की ओर बहने से रोकता है। इसके अलावा, LAMB को सीखने की दर और वार्मअप चरणों के सावधानीपूर्वक ट्यूनिंग की आवश्यकता होती है; एक अनुचित अनुसूची विचलन का कारण बन सकती है।
मेमोरी उपयोग एडम के समान है (प्रति पैरामीटर दो मोमेंट वेक्टर), इसलिए यह अधिक मेमोरी-हंगर नहीं है। बहुत बड़े मॉडल के लिए, मॉडल प्रूनिंग या ग्रेडिएंट क्लिपिंग का उपयोग LAMB के साथ किया जा सकता है, लेकिन ये ऑर्थोगोनल तकनीकें हैं।
निष्कर्ष
LAMB बड़े पैमाने पर गहन शिक्षण के लिए अनुकूलन एल्गोरिदम के टूलबॉक्स में एक आधारशिला बन गया है। विशाल मिनी-बैच के साथ प्रभावी प्रशिक्षण को सक्षम करके, इसने कई बेंचमार्क मॉडल के विकास को तेज किया है और प्रयोग की लागत को कम किया है। इसका लेयर-वाइज़ अनुकूलन सिद्धांत बाद के ऑप्टिमाइज़र डिज़ाइनों को प्रभावित किया है और वितरित प्रशिक्षण की चुनौतियों का सामना करने वाले अभ्यासकर्ताओं के लिए एक व्यावहारिक, अच्छी तरह से समझा जाने वाला समाधान बना हुआ है। जैसे-जैसे कृत्रिम बुद्धिमत्ता बढ़ती जा रही है, LAMB जैसे ऑप्टिमाइज़र विकसित होने की संभावना है, लेकिन प्रति-लेयर विश्वास और एडेप्टिव मोमेंट्स के इसके मूल विचार यहां रहने के लिए हैं।