An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

एक Mamba मॉडल की कल्पना एक विशाल, अत्यंत बुद्धिमान ऑर्केस्ट्रा के रूप में करें। इस ऑर्केस्ट्रा के हर लेयर के भीतर, 16 छोटे, एक-सुर वाले वाद्य यंत्रों (जिन्हें "मोड्स" कहा जाता है) का एक बैंक है। पुराने सोचने के तरीके में, हमने माना था कि कंडक्टर (मॉडल) पूरे गाने के लिए, चाहे संगीत कुछ भी कर रहा हो, मान लीजिए 8 सबसे अच्छे वाद्य यंत्रों का एक निश्चित सेट चुनेगा।

लेकिन यह पेपर पर्दा हटाकर एक चौंकाने वाला रहस्य उजागर करता है: कंडक्टर कोई निश्चित सेट नहीं चुनता है।

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

जादुई उपकरण: "सटीक वाद्य यंत्र" (The Exact Instrument)

लेखकों को यह कैसे पता चला? उन्होंने एक गणितीय "सटीक वाद्य यंत्र" बनाया।

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

उन्होंने इस टूल का वास्तविक मॉडल के विरुद्ध परीक्षण किया और पाया कि यह 2.3 × 10⁻⁷ के सापेक्ष त्रुटि (relative error) तक सटीक था। यह पृथ्वी से चंद्रमा तक की दूरी मापने और उसमें मानव बाल की चौड़ाई से भी कम की चूक होने जैसा है। यह कोई अनुमान नहीं है; यह एक सटीक माप है।

बड़ी खोज: "माइग्रेशन गैप" (The Migration Gap)

इस टूल का उपयोग करते हुए, उन्होंने छोटे मॉडलों (130M पैरामीटर्स) से लेकर विशाल मॉडलों (7B पैरामीटर्स, जैसे तैनात Falcon-Mamba) तक के मॉडलों का अध्ययन किया।

उन्होंने पाया कि सबसे सक्रिय लेयर्स में, एक निश्चित सेट (fixed set) में एक बदलते हुए सेट की तुलना में दोगुनी अधिक त्रुटि होती है।

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

यह प्रत्येक मॉडल में हुआ जिसका उन्होंने परीक्षण किया: Mamba-1 परिवार, 7B Falcon-Mamba, और यहाँ तक कि Mamba-2 भी।

माइग्रेशन का कारण क्या है? (क्यों?)

लेखकों ने पूछा: मॉडल का कौन सा हिस्सा यह सारा स्विचिंग कर रहा है?
एक Mamba लेयर में तीन मुख्य संकेत (signals) होते हैं:

1. राइट मैप ($B_t$): तय करता है कि किस वाद्य यंत्र को इनपुट सिग्नल मिलता है।
2. रीडआउट ($C_t$): तय करता है कि किन वाद्य यंत्रों को सुना जाता है।
3. टाइमस्टेप ($\Delta$): जिसे अक्सर "सेलेक्टिविटी" (चयनशीलता) का नॉब माना जाता है।

उन्होंने एक "फ्रोजन-सिग्नल" प्रयोग चलाया। उन्होंने यह देखने के लिए प्रत्येक सिग्नल को उसके औसत मान पर फ्रीज (स्थिर) कर दिया कि क्या माइग्रेशन रुक जाता है।

• परिणाम: जब उन्होंने राइट मैप ($B_t$) को फ्रीज किया, तो माइग्रेशन गायब हो गया। मॉडल ने वाद्य यंत्र बदलना बंद कर दिया।
• आश्चर्य: जब उन्होंने टाइमस्टेप ($\Delta$) को फ्रीज किया, तो माइग्रेशन बिल्कुल वैसा ही रहा।

निष्कर्ष: "टाइमस्टेप" सिग्नल, जिसे कई लोग चयनशीलता की कुंजी मानते थे, माइग्रेशन सिग्नल में लगभग कोई योगदान नहीं देता है। असली नायक राइट मैप ($B_t$) है। यह वह गेटकीपर है जो टोकन दर टोकन तय करता है कि किन वाद्य यंत्रों को बजने दिया जाए।

इसका परिणाम: क्या हम इसका उपयोग कर सकते हैं?

लेखकों ने इस ज्ञान का उपयोग मॉडल को छोटा (प्रून) करने के लिए किया ताकि स्थान बचाया जा सके।

• स्टैटिक प्रूनिंग (Static Pruning): एक औसत परीक्षण के आधार पर सर्वश्रेष्ठ 8 वाद्य यंत्र चुनें और उन्हें हमेशा के लिए रखें।
• इनपुट-शेड्यूल्ड प्रूनिंग (Input-Scheduled Pruning): वर्तमान वाक्य को देखें, मापें कि अभी कौन से 8 वाद्य यंत्र सक्रिय हैं, और केवल उन्हें ही रखें।

परिणाम:
आधे स्टेट बजट (16 में से केवल 8 मोड रखना) पर, इनपुट-शेड्यूल्ड विधि ने कच्चे सटीकता (raw accuracy) के मामले में पूर्ण, अन-प्रून्ड मॉडल के बराबर (और कुछ मामलों में थोड़ा बेहतर) प्रदर्शन किया।

• 130M मॉडल पर, शेड्यूल्ड विधि की पर्प्लेक्सिटी (perplexity) 11.84 थी, जबकि अन-प्रून्ड मॉडल की 12.38 थी।
• 7B Falcon-Mamba पर, शेड्यूल्ड विधि 4.44 थी, जिसने अन-प्रून्ड के 4.48 को पीछे छोड़ दिया।

हालाँकि, एक महत्वपूर्ण शर्त है: पेपर स्पष्ट रूप से कहता है कि यह "शेड्यूल्ड" विधि एक दो-पास ओरकल (two-pass oracle) है। यह पहले पूरे विंडो को तय करने के लिए एक बार पढ़ता है, फिर आउटपुट जेनरेट करने के लिए दूसरा पास चलाता है। इसका मतलब है कि यह वास्तविक दुनिया के परिनियोजन (deployment) में कंप्यूट या मेमोरी नहीं बचाता है (आपको डेटा को दो बार पढ़ना पड़ता है)।

लेखक स्पष्ट करते हैं कि यह परिणाम वास्तविक क्षमता (realizable headroom) को दर्शाता है, न कि एक तैनात दक्षता जीत (deployed efficiency win)। यह सिद्ध करता है कि एक छोटा, सुपर-एफिशिएंट मॉडल विशाल मॉडल के बराबर होने की क्षमता रखता है, लेकिन केवल तभी जब हम एक सस्ता, तेज़ प्रेडिक्टर बना सकें जो महंगे पहले पास की आवश्यकता के बिना सही वाद्य यंत्रों का अनुमान लगा सके। वर्तमान विधि केवल यह दिखाती है कि क्या संभव है।

उन्होंने क्या खारिज किया

पेपर बहुत स्पष्ट है कि क्या काम नहीं करता है:

• स्टैटिक रैंकिंग्स (Static Rankings): वे तरीके जो औसत गतिविधि या "हैंकेल एनर्जी" (जैसे GHOST या LAST) के आधार पर मोड का एक निश्चित सेट चुनते हैं, वे काफी खराब हैं। वे चलते हुए लक्ष्य के साथ तालमेल नहीं बिठा पाते।
• टाइमस्टेप एक्टिविटी (Timestep Activity): टाइमस्टेप की सक्रियता के आधार पर प्रूनिंग करना एक बेकार रास्ता है क्योंकि टाइमस्टेप माइग्रेशन को संचालित नहीं करता है।
• सरल प्रेडिक्टर्स (Simple Predictors): उन्होंने केवल वाक्य के पहले आधे हिस्से या "डोमेन" (जैसे कोड बनाम गद्य) के आधार पर मास्क की भविष्यवाणी करने की कोशिश की। ये सस्ते तरीके केवल संभावित लाभ का 2-6% ही प्राप्त कर सके। "महत्वपूर्ण" वाद्य यंत्रों का सेट इतनी तेज़ी से बदलता है (कुछ सौ टोकन के भीतर) कि पूर्ण लाभ प्राप्त करने के लिए आपको स्कोर किए जा रहे विशिष्ट टोकन को मापना ही होगा।

निचोड़

पेपर यह सिद्ध करता है कि प्रशिक्षित चयनात्मक स्टेट-स्पेस मॉडल गतिशील, जीवित चीजें हैं जो इनपुट के आधार पर अपने आंतरिक संसाधनों को लगातार पुनर्वितरित करते हैं। "राइट मैप" इस माइग्रेशन का कंडक्टर है। हालाँकि हम अभी एक ऐसा मॉडल तैनात नहीं कर सकते जो बिना "दो-पास" लागत के वास्तविक समय में अपने वाद्य यंत्रों को बदल सके, फिर भी यह शोध हमें क्षेत्र का एक सटीक मानचित्र देता है। यह हमें दिखाता है कि वर्तमान "स्टैटिक" प्रूनिंग विधियाँ प्रदर्शन का एक बड़ा हिस्सा छोड़ रही हैं, और भविष्य का मार्ग ऐसे शेड्यूलर बनाने में निहित है जो इन माइग्रेशन का मौके पर ही अनुमान लगा सकें।