RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale

तकनीकी सारांश: RADLADS

समस्या विवरण

जबकि लीनियर अटेंशन वेरिएंट (जैसे RWKV और Mamba) पारंपरिक सॉफ्टमैक्स अटेंशन ट्रांसफॉर्मर्स की तुलना में महत्वपूर्ण लाभ प्रदान करते हैं—विशेष रूप से की-वैल्यू (KV) कैश को टालकर प्रति टोकन $O(1)$ इन्फरेंस समय और निरंतर मेमोरी उपयोग—बड़े पैमाने के लीनियर मॉडल को शून्य से प्रशिक्षित करना अत्यधिक महंगा बना हुआ है। अत्याधुनिक (SoTA) ट्रांसफॉर्मर मॉडल्स को अक्सर दस ट्रिलियन से अधिक टोकन पर प्रशिक्षित करने की आवश्यकता होती है, जो एक उच्च लागत है जो अधिकांश संगठनों के लिए सुलभ नहीं है। प्री-ट्रेंड सॉफ्टमैक्स ट्रांसफॉर्मर्स को लीनियर या रिकरेंट मॉडल्स में बदलने के मौजूदा तरीके (जैसे T2R, SUPRA, LoLCats, MOHAWK) ऐतिहासिक रूप से भारी टोकन काउंट (20B से 100B+ टोकन तक) की मांग करते हैं या खराब डाउनस्ट्रीम प्रदर्शन (विशेष रूप से MMLU जैसे बेंचमार्क पर) का परिणाम देते हैं। इसके अलावा, कई रूपांतरण प्रयास हाइब्रिड आर्किटेक्चर (कुछ सॉफ्टमैक्स अटेंशन को बनाए रखना) पर निर्भर करते हैं या टीचर मॉडल्स की गुणवत्ता से मेल खाने में विफल रहते हैं।

कार्यप्रणाली

लेखक RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale) प्रस्तुत करते हैं, जो एक प्रोटोकॉल है जिसे न्यूनतम कम्प्यूटेशनल लागत के साथ सॉफ्टमैक्स अटेंशन ट्रांसफॉर्मर्स को लीनियर अटेंशन डीकोडर मॉडल में बदलने के लिए डिज़ाइन किया गया है। इस प्रक्रिया में तीन प्राथमिक चरण और विशिष्ट आर्किटेक्चरल अनुकूलन शामिल हैं:

1. RADLADS प्रोटोकॉल

रूपांतरण एक तीन-चरणीय प्रक्रिया है जिसके लिए केवल 350–700 मिलियन टोकन (टीचर के प्री-ट्रेनिंग डेटा के 0.005% से भी कम) की आवश्यकता होती है:

• सेटअप (अटेंशन वेट्स ट्रांसफर): टीचर मॉडल से अटेंशन-संबंधित वेट्स ($W_q, W_k, W_v, W_o$) सीधे स्टूडेंट मॉडल के सीक्वेंस मिक्सिंग लेयर्स में ट्रांसफर किए जाते हैं। अन्य वेट्स को मानक प्री-ट्रेनिंग विधियों के माध्यम से इनिशियलाइज़ किया जाता है या उन्हें बिना तत्काल प्रभाव के टीचर के व्यवहार की नकल करने के लिए सेट किया जाता है।
• चरण 1: अटेंशन हिडन स्टेट अलाइनमेंट: स्टूडेंट मॉडल के सीक्वेंस मिक्सिंग लेयर्स को संबंधित टीचर अटेंशन लेयर्स के हिडन स्टेट आउटपुट्स को अनुमानित करने के लिए प्रशिक्षित किया जाता है। यह एक L2 डिस्टेंस (या MSE) ऑब्जेक्टिव का उपयोग करके किया जाता है। लेखकों ने पाया कि एक गेटेड लीनियर अटेंशन कर्नेल (ऑफ-बाय-वन डिके और बोनस टर्म्स को हटाकर) का उपयोग करने से मानक RWKV-6 की तुलना में टीचर के हिडन स्टेट्स के साथ बेहतर फिट प्राप्त होता है।
• चरण 2: नॉलेज डिस्टिलेशन: संपूर्ण स्टूडेंट मॉडल को टीचर के आउटपुट लॉजिट्स को अनुमानित करने के लिए KL डाइवर्जेंस लॉस का उपयोग करके प्रशिक्षित किया जाता है। लेखक परिकल्पना करते हैं कि तथ्यात्मक ज्ञान मुख्य रूप से टीचर के MLPs और एम्बेडिंग्स में रहता है; इसलिए, कैटास्ट्रोफिक फॉरगेटिंग को रोकने के लिए इन घटकों के लिए लर्निंग रेट को कम या स्थिर रखा जाता है, जबकि सीक्वेंस मिक्सर को अधिक आक्रामक रूप से प्रशिक्षित किया जाता है।
• चरण 3: कॉन्टेक्स्ट लेंथ एक्सटेंशन (वैकल्पिक): मॉडल को लंबी-संदर्भ क्षमताओं को बढ़ाने के लिए बिना किसी टीचर मॉडल के मानक क्रॉस-एन्ट्रॉपी लॉस का उपयोग करके लंबे सीक्वेंस (16k टोकन तक) पर फाइन-ट्यून किया जाता है। वैकल्पिक रूप से, लेखक चरण 2a का प्रस्ताव करते हैं, जो डिस्टिलेशन चरण के दौरान ही सीक्वेंस लेंथ को 4096 तक बढ़ाता है, जिससे संभावित रूप से एक अलग चरण 3 की आवश्यकता समाप्त हो जाती है।

2. नए आर्किटेक्चर

लेखकों ने पहचान की कि मानक RWKV आर्किटेक्चर हमेशा रूपांतरण के लिए इष्टतम नहीं थे। उन्होंने दो नए वेरिएंट पेश किए:

• RAD-RWKV6 ("RADFinch"): एक गेटेड लीनियर अटेंशन कर्नेल और स्टेट बैलेंसिंग तकनीकों का उपयोग करने वाला RWKV-6 का एक संशोधन, जो स्थिरता और फिट में सुधार करता है।
• RAD-RWKV7 ("RADGoose"): एक RWKV-7 का संशोधन जिसमें "टोकनशिफ्ट" तंत्र को हटा दिया गया है (जिसने इस संदर्भ में कोई लाभ नहीं दिया) और सीधे रोटरी पोजीशनल एम्बेडिंग्स (RoPE) लागू की गई है। इस आर्किटेक्चर ने अनमॉडिफाइड RWKV-7 की तुलना में तेज़ अभिसरण (convergence) और कम डिस्टिलेशन लॉस प्रदर्शित किया।

3. हाइपरपैरामीटर्स और डेटा

• डेटासेट: लेखकों ने सभी रूपांतरण चरणों के लिए DCLM (DataComp-LM) को चुना, जिसे Qwen मॉडल्स को रूपांतरित करने के लिए FineWeb या FineWeb-Edu की तुलना में बेहतर पाया गया।
• लर्निंग रेट्स: चरण 1 में एक कोसाइन एनीलिंग शेड्यूल का उपयोग किया जाता है, जो उच्च ($10^{-3}$) से शुरू होकर टीचर के अंतिम प्री-ट्रेनिंग लर्निंग रेट ($10^{-5}$) के करीब समाप्त होता है। चरण 2 और 3 एक फ्लैट लर्निंग रेट का उपयोग करते हैं।

मुख्य योगदान

1. RADLADS डिस्टिलेशन रेसिपी: एक विस्तृत, चरण-दर-चरण प्रोटोकॉल जिसमें विशिष्ट हाइपरपैरामीटर्स, टोकन काउंट और डेटासेट विकल्प शामिल हैं जो न्यूनतम डेटा के साथ उच्च-गुणवत्ता वाले रूपांतरण को सक्षम बनाता है।
2. नए आर्किटेक्चर: RAD-RWKV6 और RAD-RWKV7 का परिचय, जो रूपांतरण प्रक्रिया के लिए अनुकूलित हैं, जिससे उनके अनमॉडिफाइड पूर्ववर्तियों की तुलना में तेज़ इन्फरेंस और बेहतर टीचर अलाइनमेंट प्राप्त होता है।
3. बड़े पैमाने के मॉडल: 7B, 32B, और 72B पैरामीटर स्केल पर लोकप्रिय ओपन-सोर्स Qwen2.5 मॉडल्स का सफल रूपांतरण।
4. ओपन सोर्स रिलीज़: कोड और रूपांतरित मॉडल्स (QRWKV6/7-7B/32B/72B) को Apache 2.0 लाइसेंस (72B मॉडल के लिए Qwen लाइसेंस प्रतिबंधों के साथ) के तहत जारी करना, जिससे अन्य लोग इस प्रक्रिया को दोहरा सकें।

परिणाम

रूपांतरित मॉडल अपने आकार के शुद्ध रिकरेंट मॉडल्स के बीच अत्याधुनिक प्रदर्शन प्राप्त करते हैं:

• दक्षता: एक 72B मॉडल को रूपांतरित करने की लागत $2,000 USD से कम है और इसमें केवल ~700M टोकन लगते हैं।
• प्रदर्शन: मानक बेंचमार्क (Lambada, MMLU, ARC, आदि) पर, RADLADS मॉडल्स लगातार अन्य रूपांतरण विधियों (जैसे SUPRA, LoLCats, MOHAWK, ARWKV) को पछाड़ते हैं।
  • QRWKV7-7B-Instruct अपने टीचर की तुलना में 92.4% का सापेक्ष MMLU स्कोर प्राप्त करता है, जो अन्य 7B रूपांतरणों से काफी बेहतर है।
  • QRWKV6-72B-Instruct 89.9% का सापेक्ष MMLU स्कोर प्राप्त करता है, जो उस स्केल पर शुद्ध RNN लैंग्वेज मॉडल्स के लिए एक नया SoTA स्थापित करता है।
• इन्फरेंस स्पीड: लीनियर अटेंशन मैकेनिज्म के कारण, रूपांतरित मॉडल कॉन्टेक्स्ट लेंथ बढ़ने के साथ टीचर मॉडल्स की तुलना में महत्वपूर्ण गति वृद्धि दिखाते हैं। उदाहरण के लिए, 8k इनपुट टोकन और 256 आउटपुट टोकन पर, 32B QRWKV7 मॉडल अपने टीचर Qwen3-32B की तुलना में 1.61x तेज़ है; 6k इनपुट और 2k आउटपुट पर, स्पीडअप 3.41x तक पहुँच जाता है।

सीमाएं और दावे

पेपर विनम्रतापूर्वक कुछ सीमाओं को स्वीकार करता है:

• रीजनिंग और लॉन्ग कॉन्टेक्स्ट: हालांकि मानक बेंचमार्क पर प्रदर्शन मजबूत है, लेकिन मॉडल्स जटिल रीजनिंग कार्यों (जैसे Minerva Math) और बहुत लंबे-संदर्भ वाले कार्यों (जैसे RULER) में सीमाएं दिखाते हैं, जहाँ वे टीचर मॉडल्स की तरह अतिरिक्त आउटपुट टोकन्स के साथ प्रभावी ढंग से सुधार नहीं करते हैं।
• आर्किटेक्चर संवेदनशीलता: प्रत्येक नए आर्किटेक्चर डिज़ाइन के लिए यह सुनिश्चित करने के लिए सूक्ष्म परीक्षण की आवश्यकता होती है कि वे RADLADS प्रोटोकॉल के साथ संगत हों। उदाहरण के लिए, GroupNorm/LayerNorm, जो कुछ लीनियर अटेंशन वेरिएंट्स में सामान्य है, 14B+ स्केल पर प्रशिक्षण अस्थिरता का कारण बना और इसे प्री-स्केलिंग तकनीकों के साथ बदलना पड़ा।
• डेटासेट अलाइनमेंट: लेखक नोट करते हैं कि रिपीटिटिव लूपिंग व्यवहार को रोकने के लिए रीजनिंग डेटासेट्स को टीचर मॉडल के डिस्ट्रीब्यूशन के साथ अधिक निकटता से संरेखित करने की आवश्यकता हो सकती है, एक चुनौती जिसे उन्होंने DCLM को OpenThoughts के साथ मिलाकर आंशिक रूप से संबोधित किया है।

महत्व

यह पेपर दावा करता है कि RADLADS बड़े पैमाने के लीनियर अटेंशन मॉडल्स तक पहुंच को लोकतांत्रिक बनाने के लिए एक लागत प्रभावी मार्ग प्रदान करता है। प्रशिक्षण टोकन आवश्यकता को ट्रिलियनों से सैकड़ों मिलियनों तक कम करके, यह शोधकर्ताओं और छोटे संगठनों को अत्यधिक लागत के बिना नए प्रकार के RNN आर्किटेक्चर को बड़े पैमाने पर परीक्षण करने, प्रशिक्षित करने और तैनात करने में सक्षम बनाता है। लेखक इसे कंप्रेसिव स्टेट अटेंशन वेरिएंट्स के विकास को गति देने वाले एक उपकरण के रूप में देखते हैं।