तकनीकी सारांश: एसोसिएटिव रिकरेंट मेमोरी के माध्यम से LLM कॉन्टेक्स्ट का विस्तार

समस्या विवरण

लार्ज लैंग्वेज मॉडल्स (LLMs) को तेजी से ऐसे इनपुट्स प्रोसेस करने की आवश्यकता है जो तकनीकी रिपोर्ट विश्लेषण, सॉफ्टवेयर डेवलपमेंट और मल्टी-डॉक्यूमेंट रीजनिंग जैसे कार्यों के लिए सैकड़ों हजार या लाखों टोकन तक फैले हों। हालांकि, मानक ट्रांसफॉर्मर आर्किटेक्चर एक मौलिक बाधा का सामना करते हैं: सेल्फ-अटेंशन की कम्प्यूटेशनल और मेमोरी लागत अनुक्रम लंबाई (sequence length) के साथ द्विघात रूप से ($O(N^2)$) बढ़ती है। इसके अलावा, प्रदर्शन अक्सर घट जाता है जब कॉन्टेक्स्ट की लंबाई बढ़ती है, भले ही वह मॉडल की नाममात्र की विंडो के भीतर ही क्यों न हो। जबकि रिकरेंट आर्किटेक्चर (जैसे Mamba, RWKW) रैखिक स्केलिंग (linear scaling) प्रदान करते हैं, उन्हें आमतौर पर शून्य से प्रशिक्षित करने की आवश्यकता होती है, जिससे मौजूदा प्री-ट्रेंड LLMs का लाभ उठाना कठिन हो जाता है, और वे अक्सर ट्रांसफॉर्मर्स की तुलना में जटिल एल्गोरिद्मिक कार्यों और इंस्ट्रक्शन फॉलोइंग में संघर्ष करते हैं।

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

लेखक एसोसिएटिव रिकरेंट मेमोरी ट्रांसफॉर्मर (ARMT) को कॉन्टेक्स्ट की लंबाई बढ़ाने के लिए एक व्यावहारिक समाधान के रूप में प्रस्तावित करते हैं, जो निरंतर मेमोरी स्केलिंग और दक्षता बनाए रखता है। ARMT एक प्री-ट्रेंड बेस LLM के चारों ओर एक 'रैपर' के रूप में कार्य करता है, जो इनपुट के सेगमेंट-वार प्रोसेसिंग को सक्षम बनाता है।

मुख्य आर्किटेक्चर

ARMT लंबे-कॉन्टेक्स्ट इनपुट्स को निश्चित लंबाई के गैर-ओवरलैपिंग सेगमेंट्स में विभाजित करता है। प्रत्येक सेगमेंट के भीतर, मॉडल फुल सेल्फ-अटेंशन (शॉर्ट-टर्म/वर्किंग मेमोरी) का उपयोग करता है। महत्वपूर्ण रूप से, यह सूचना को सेगमेंट्स के बीच प्रसारित करने के लिए एक लेयरवाइज एसोसिएटिव मेमोरी मॉड्यूल (layerwise associative memory module) पेश करता है (लॉन्ग-टर्म मेमोरी)। यह तंत्र तीन चरणों में संचालित होता है:

1. मेमोरी एक्सट्रैक्शन (Memory Extraction): प्रत्येक ट्रांसफॉर्मर लेयर इनपुट सेगमेंट को मेमोरी एम्बेडिंग्स में संकुचित (compress) करती है।
2. मेमोरी कंसोलिडेशन (Memory Consolidation): इन एम्बेडिंग्स को की-वैल्यू पेयर्स (key-value pairs) के रूप में प्रति-लेयर एसोसिएटिव मैट्रिक्स में समेकित किया जाता है।
3. एसोसिएशन (Association): आगामी सेगमेंट्स के एम्बेडिंग्स को क्वेरी वेक्टर्स में बदला जाता है और पिछले सेगमेंट्स से प्रासंगिक जानकारी प्राप्त करने के लिए एसोसिएटिव मैट्रिक्स के साथ गुणा किया जाता है।

ट्रेनिंग रेसिपी (Training Recipe)

पेपर प्री-ट्रेंड LLMs को ARMT के अनुकूल बनाने के लिए एक व्यापक प्रशिक्षण रणनीति का वर्णन करता है:

• कंटीन्यूड प्री-ट्रेनिंग (Continued Pre-training): अन-इनिशियलाइज्ड एसोसिएटिव मेमोरी पैरामीटर्स को लंबे कॉन्टेक्स्ट (जैसे FineWeb-Edu से 19B टोकन) पर अनसुपरवाइज्ड लैंग्वेज मॉडलिंग के माध्यम से इनिशियलाइज किया जाता है ताकि टास्क-स्पेसिफिक फाइन-ट्यूनिंग से पहले प्रभावी मेमोरी प्रोपेगेशन सीखा जा सके।
• करिकुलम लर्निंग (Curriculum Learning): लॉन्ग-रेंज डिपेंडेंसीज को शून्य से सीखने की कठिनाई को दूर करने के लिए, मॉडल को सेगमेंट्स की संख्या को धीरे-धीरे बढ़ाकर (जैसे 2 से 4 से 8 तक) और लर्निंग रेट को कम करके (annealing) फाइन-ट्यून किया जाता है।
• सिंथेटिक डेटा जनरेशन (Synthetic Data Generation): लंबे-कॉन्टेक्स्ट परिदृश्यों में डेटा की कमी को दूर करने के लिए, लेखक लंबे दस्तावेजों के छोटे पैसेजेस को जोड़कर और प्रत्येक के लिए QA पेयर्स जेनरेट करके सिंथेटिक ट्रेनिंग इंस्टेंस तैयार करते हैं, जिससे विशिष्ट कॉन्टेक्स्ट लेंथ के बड़े बिन (bins) बनते हैं।
• लेयर प्रूनिंग और सिलेक्शन (Layer Pruning and Selection): अध्ययन इस बात की जांच करता है कि क्या हर लेयर में एसोसिएटिव मेमोरी की आवश्यकता है। यह केवल कुछ चुनिकी गई लेयर्स (जैसे विशिष्ट मिडिल और फाइनल लेयर्स) में एसोसिएटिव मेमोरी रखने का प्रस्ताव देता है, जिससे प्रदर्शन में महत्वपूर्ण गिरावट के बिना ट्रेन करने योग्य पैरामीटर्स और कम्प्यूटेशनल लागत कम हो जाती है।

मुख्य योगदान

1. डोमेन-विशिष्ट डेटासेट: दो नए डेटासेट का निर्माण, ManyTypes-long (MT) कोड में वेरिएबल टाइप प्रेडिक्शन के लिए और GovReport-long (GR) लंबे दस्तावेज़ों के प्रश्न-उत्तर के लिए, जो वास्तविक, नैरो-डोमेन वर्कलोड का मूल्यांकन करने के लिए डिज़ाइन किए गए हैं।
2. ट्रेनिंग रेसिपी: ARMT का उपयोग करके LLM कॉन्टेक्स्ट को विस्तारित करने के लिए एक नया फ्रेमवर्क, जो कंटीन्यूड प्री-ट्रेनिंग, सिंथेटिक डेटा जनरेशन, करिकुलम लर्निंग और सिलेक्टिव लेयर इंटीग्रेशन को जोड़ता है।
3. अनुभवजन्य सत्यापन (Empirical Validation): एक विस्तृत प्रयोगात्मक अध्ययन जो प्रदर्शित करता है कि ARMT-ऑगमेंटेड मॉडल:
   • मूल कॉन्टेक्स्ट लिमिट्स से कहीं अधिक (64k टोकन तक) इनपुट्स को बिना किसी प्रदर्शन गिरावट के प्रोसेस करते हैं।
   • बेस मॉडल्स की तुलना में आउट-ऑफ-डिस्ट्रीब्यूशन (OOD) कॉन्टेक्स्ट लेंथ के प्रति बेहतर सामान्यीकरण (generalization) प्रदर्शित करते हैं।
   • फुल-अटेंशन मॉडल्स की तुलना में बेसलाइन प्रदर्शन को बनाए रखते हुए लगभग 30% कम FLOPs की आवश्यकता होती है।

प्रयोगात्मक परिणाम

लेखकों ने MT और GR डेटासेट पर Gemma-3-1B-IT और SmolLM-2-360M-IT बैकबोन का उपयोग करके ARMT का मूल्यांकन किया।

• प्रदर्शन (Performance): ARMT मॉडल्स ने 65k टोकन तक की कॉन्टेक्स्ट लेंथ में स्थिर प्रदर्शन बनाए रखा। इसके विपरीत, बेस मॉडल्स (भले ही उन्हें फाइन-ट्यून किया गया हो) अपनी नेटिव कॉन्टेक्स्ट विंडोज (जैसे 8k या 32k) के बाद प्रदर्शन में भारी गिरावट दिखाते हैं। ARMT ने लॉन्ग-OOD रिजीम (32k–65k) में बेस मॉडल्स को काफी पीछे छोड़ दिया।
• दक्षता (Efficiency): ARMT ने कॉन्टेक्स्ट की लंबाई के बावजूद स्थिर GPU मेमोरी उपयोग प्रदर्शित किया, जबकि बेस मॉडल्स का मेमोरी उपयोग रैखिक रूप से बढ़ता गया। 32k टोकन के सीक्वेंस के लिए, ARMT ने समान मेमोरी बजट के तहत बैच साइज को 4 गुना बढ़ा दिया (32 बनाम 8)।
• FLOPs में कमी: थ्योरिटिकल एनालिसिस और एम्पिरिकल इन्फरेंस टाइम मेजरमेंट्स ने ग्लोबल अटेंशन FLOPs में $T/S$ (सीक्वेंस लेंथ विभाजित सेगमेंट साइज) के कारक से कमी की पुष्टि की, जिससे फुल-अटेंशन मॉडल्स की तुलना में कुल FLOPs में लगभग ~30% की कमी आई।
• एब्लेशन स्टडीज (Ablation Studies):
  • लेयर प्रूनिंग: केवल ~20% लेयर्स (विशेष रूप से पूर्व-चयनित मिडिल और फाइनल लेयर्स) में एसोसिएटिव मेमोरी वाले मॉडल्स ने फुल ARMT मॉडल्स के बराबर या बेहतर प्रदर्शन हासिल किया।
  • प्री-ट्रेनिंग: एसोसिएटिव मेमोरी को इनिशियलाइज करने के लिए कंटीन्यूड प्री-ट्रेनिंग को आवश्यक पाया गया, जिसने इन-डोमेन और OOD दोनों प्रदर्शन में सुधार किया।
  • बेसलाइन्स: ARMT ने अन्य लॉन्ग-कॉन्टेक्स्ट बेसलाइन्स, जिनमें Mamba-2, DeltaNet, और xLSTM शामिल हैं, विशेष रूप से लॉन्ग-कॉन्टेक्स्ट जनरलाइजेशन में बेहतर प्रदर्शन किया, जबकि इसे शून्य से प्रशिक्षित मॉडल्स की तुलना में कम व्यापक प्री-ट्रेनिंग की आवश्यकता थी।

महत्व और दावे

यह पेपर ARMT को छोटे से मध्यम आकार के LLMs (1B पैरामीटर्स तक) में लॉन्ग-कॉन्टेक्स्ट प्रोसेसिंग को सक्षम करने के लिए एक व्यावहारिक, कंप्यूट-एफिशिएंट दृष्टिकोण के रूप में स्थापित करता है। लेखक दावा करते हैं कि यह दृष्टिकोण ट्रांसफॉर्मर्स के मजबूत शॉर्ट-कॉन्टेक्स्ट प्रदर्शन और रिकरेंट मॉडल्स की रैखिक स्केलिंग के बीच के अंतर को पाटता है।

महत्व के संबंध में प्रमुख दावे शामिल हैं:

• प्राइवेसी और लोकल डिप्लॉयमेंट: छोटे मॉडल्स में कुशल लॉन्ग-कॉन्टेक्स्ट प्रोसेसिंग को सक्षम करके, ARMT उन प्राइवेसी-प्रिजर्विंग एप्लिकेशन्स को सुगम बनाता है जो रिमोट, बड़े-पैमाने के API-आधारित LLMs पर निर्भर नहीं हैं।
• स्केलेबिलिटी: यह विधि स्थिर मेमोरी के साथ अनिश्चित रूप से लंबे कॉन्टेक्स्ट को प्रोसेस करने की अनुमति देती है, जो वास्तविक दुनिया के डॉक्यूमेंट विश्लेषण के लिए एक महत्वपूर्ण आवश्यकता है।
• दक्षता: निरंतर मेमोरी स्केलिंग और कम FLOPs का संयोजन ARMT को संसाधन-बाधित वातावरण के लिए एक व्यवहार्य विकल्प बनाता है।
• जनरलाइजेशन: प्रस्तावित ट्रेनिंग रेसिपी "लॉस्ट इन द मिडल" घटना और मानक ट्रांसफॉर्मर्स में देखी जाने वाली प्रदर्शन गिरावट को प्रभावी ढंग से संबोधित करती है जब वे लंबे अनुक्रमों को संभालते हैं।

लेखक सीमाओं को स्वीकार करते हैं, यह नोट करते हुए कि प्रयोग 1B पैरामीटर्स तक के मॉडल्स और विशिष्ट कार्यों (कोड और डॉक्यूमेंट QA) तक सीमित थे, और कि एसोसिएटिव मेमोरी के ट्रांसफॉर्मर रिप्रेजेंटेशन्स के साथ इंटरेक्शन के अंतर्निहित तंत्र अभी भी आंशिक रूप से समझे जाने बाकी हैं। हालांकि, उनका तर्क है कि छोटे मॉडल्स में कॉन्टेक्स्ट को विस्तारित करने की प्रदर्शित क्षमता व्यावहारिक, लोकल लॉन्ग-कॉन्टेक्स्ट AI की दिशा में एक महत्वपूर्ण कदम है।