हाल ही में, Mamba टीम का शोध ध्यान आकर्षित कर रहा है: कॉर्नेल और प्रिंस्टन जैसे विश्वविद्यालयों के शोधकर्ताओं ने Llama, एक बड़े Transformer मॉडल को "डिस्टिल" करके Mamba में सफलतापूर्वक परिवर्तित किया है, और एक नई प्रकार की डिकोडिंग एल्गोरिदम डिजाइन की है, जिससे मॉडल की इनफेरेंस गति में काफी वृद्धि हुई है।
शोधकर्ताओं का लक्ष्य Llama को Mamba में बदलना है। ऐसा क्यों किया जा रहा है? क्योंकि एक बड़े मॉडल को शून्य से प्रशिक्षित करना बहुत महंगा है, और Mamba को उसके जन्म के बाद से व्यापक ध्यान मिला है, लेकिन वास्तव में बहुत कम टीमें बड़े पैमाने पर Mamba मॉडल को स्वयं प्रशिक्षित करती हैं। हालांकि बाजार में कुछ प्रसिद्ध रूपांतर हैं, जैसे AI21 का Jamba और NVIDIA का Hybrid Mamba2, लेकिन कई सफल Transformer मॉडलों में समृद्ध ज्ञान छिपा है। यदि हम इन ज्ञान को लॉक कर सकें, जबकि Transformer को Mamba के लिए ट्यून कर सकें, तो समस्या का समाधान हो जाएगा।
शोध टीम ने प्रोग्रेसिव डिस्टिलेशन, सुपरवाइज्ड फाइन-ट्यूनिंग और डायरेक्ट प्रेफरेंस ऑप्टिमाइजेशन जैसी कई विधियों को मिलाकर इस लक्ष्य को सफलतापूर्वक प्राप्त किया। ध्यान देने योग्य बात यह है कि प्रदर्शन को बनाए रखते हुए गति भी अत्यंत महत्वपूर्ण है। Mamba की लंबी अनुक्रमों में इनफेरेंस में स्पष्ट रूप से बढ़त है, जबकि Transformer के पास भी इनफेरेंस त्वरक समाधान हैं, जैसे कि अनुमानित डिकोडिंग। चूंकि Mamba की अद्वितीय संरचना इन समाधानों को सीधे लागू नहीं कर सकती, शोधकर्ताओं ने एक पूरी तरह से नई एल्गोरिदम को डिजाइन किया और Mamba आधारित अनुमानित डिकोडिंग को लागू करने के लिए हार्डवेयर विशेषताओं का उपयोग किया।
अंततः, शोधकर्ताओं ने Zephyr-7B और Llama-38B को सफलतापूर्वक एक रैखिक RNN मॉडल में परिवर्तित किया, और प्रदर्शन डिस्टिलेशन से पहले के मानक मॉडल के बराबर था। संपूर्ण प्रशिक्षण प्रक्रिया में केवल 20B टोकन का उपयोग किया गया, और परिणाम 1.2T टोकन से प्रारंभिक प्रशिक्षण के Mamba7B मॉडल और 3.5T टोकन से प्रशिक्षित NVIDIA Hybrid Mamba2 मॉडल के साथ तुलना में लगभग समान था।
तकनीकी विवरणों के संदर्भ में, रैखिक RNN और रैखिक ध्यान एक-दूसरे से जुड़े हुए हैं, इसलिए शोधकर्ताओं ने ध्यान तंत्र में प्रक्षिप्ति मैट्रिक्स को सीधे पुनः उपयोग करने में सक्षम थे, और मॉडल निर्माण को पैरामीटर प्रारंभिककरण के माध्यम से पूरा किया। इसके अतिरिक्त, शोध टीम ने Transformer में MLP परत के पैरामीटर को फ्रीज किया और धीरे-धीरे ध्यान सिर को रैखिक RNN परत (यानी Mamba) से बदल दिया, और क्रॉस-हेड साझा की गई कुंजी और मान के समूह प्रश्न ध्यान को संभाला।
डिस्टिलेशन प्रक्रिया में, ध्यान परत को क्रमिक रूप से बदलने की रणनीति अपनाई गई। सुपरवाइज्ड फाइन-ट्यूनिंग में दो मुख्य विधियाँ शामिल हैं: एक वर्ड-लेवल KL डाइवर्जेंस पर आधारित है, और दूसरी अनुक्रम-स्तरीय ज्ञान डिस्टिलेशन है। उपयोगकर्ता की प्राथमिकताओं के ट्यूनिंग चरण के लिए, टीम ने डायरेक्ट प्रेफरेंस ऑप्टिमाइजेशन (DPO) विधि का उपयोग किया, जो शिक्षक मॉडल के आउटपुट के साथ तुलना करके सुनिश्चित करता है कि मॉडल सामग्री उत्पन्न करते समय उपयोगकर्ता की अपेक्षाओं के अनुसार बेहतर हो।
इसके बाद, शोधकर्ताओं ने Transformer के अनुमानित डिकोडिंग को Mamba मॉडल में लागू करना शुरू किया। अनुमानित डिकोडिंग को सरलता से समझा जा सकता है कि एक छोटे मॉडल का उपयोग करके कई आउटपुट उत्पन्न करना, और फिर बड़े मॉडल का उपयोग करके इन आउटपुट को सत्यापित करना। छोटा मॉडल तेजी से चलता है, और कई आउटपुट वेक्टर जल्दी उत्पन्न कर सकता है, जबकि बड़ा मॉडल इन आउटपुट की सटीकता का मूल्यांकन करता है, जिससे समग्र इनफेरेंस गति बढ़ती है।
इस प्रक्रिया को लागू करने के लिए, शोधकर्ताओं ने एक एल्गोरिदम डिज़ाइन किया, जिसमें हर बार छोटे मॉडल द्वारा K ड्राफ्ट आउटपुट उत्पन्न किए जाते हैं, और फिर बड़े मॉडल द्वारा सत्यापन करके अंतिम आउटपुट और मध्यवर्ती स्थिति के कैश लौटाए जाते हैं। यह विधि GPU पर बहुत अच्छे परिणाम देती है, Mamba2.8B ने 1.5 गुना इनफेरेंस त्वरक प्राप्त किया, और स्वीकृति दर 60% तक पहुँच गई। हालांकि विभिन्न आर्किटेक्चर वाले GPU पर परिणाम भिन्न होते हैं, शोध टीम ने कोर फ्यूजन और कार्यान्वयन विधियों को समायोजित करके आगे की अनुकूलन की, अंततः आदर्श त्वरक प्रभाव प्राप्त किया।
प्रयोगात्मक चरण में, शोधकर्ताओं ने Zephyr-7B और Llama-3Instruct8B का उपयोग करके तीन चरणों में डिस्टिलेशन प्रशिक्षण किया, और अंततः केवल 8 कार्ड 80G A100 पर 3 से 4 दिनों तक चलाकर शोध परिणामों को सफलतापूर्वक पुनः उत्पन्न किया। यह अध्ययन न केवल Mamba और Llama के बीच परिवर्तन के मार्ग को दर्शाता है, बल्कि भविष्य के मॉडलों की इनफेरेंस गति और प्रदर्शन में सुधार के लिए नए विचार भी प्रदान करता है।
पेपर का पता: https://arxiv.org/pdf/2408.15237
