💬 NLP

RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference

यह शोध पत्र RAEE का प्रस्ताव करता है, जो एक सुदृढ़ रिट्रीवल-ऑगमेंटेड फ्रेमवर्क है जो समान डेटा की एग्जिट जानकारी का लाभ उठाने के लिए अर्ली एग्जिट को एक डिस्ट्रीब्यूशन प्रेडिक्शन समस्या के रूप में मॉडल करता है ताकि एडेप्टिव लेयर टर्मिनेशन को निर्देशित किया जा सके, जिससे कई कार्यों में ज़ीरो-शॉट प्रदर्शन को बनाए रखते हुए या बढ़ाते हुए LLM इन्फरेंस को त्वरित किया जा सके।

मूल लेखक: Lianming Huang, Shangyu Wu, Yufei Cui, Ying Xiong, Haibo Hu, Xue Liu, Tei-Wei Kuo, Nan Guan, Chun Jason Xue

प्रकाशित 2026-03-05

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Lianming Huang, Shangyu Wu, Yufei Cui, Ying Xiong, Haibo Hu, Xue Liu, Tei-Wei Kuo, Nan Guan, Chun Jason Xue

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप एक अत्यधिक शिक्षित विशेषज्ञ के रूप में काम कर रहे हैं जिन्हें एक लंबी, जटिल पुस्तक पढ़ने और उसके बारे में सवालों के जवाब देने के लिए काम पर रखा गया है। आमतौर पर, सटीक उत्तर पाने के लिए, आप कवर से कवर तक पूरी किताब पढ़ते हैं, चाहे सवाल कितना भी सरल क्यों न हो। वर्तमान लार्ज लैंग्वेज मॉडल्स (LLMs) इसी तरह काम करते हैं: वे हर एक इनपुट के लिए अपने "मस्तिष्क" की हर एक परत को प्रोसेस करते हैं, जिसमें बहुत समय और ऊर्जा लगती है।

समस्या:
कभी-कभी, उत्तर पहले कुछ पन्नों के बाद ही स्पष्ट हो जाता है। लेकिन क्योंकि विशेषज्ञ को पूरी किताब पढ़ने के लिए प्रोग्राम किया गया है, इसलिए वे आसान सवालों पर भी समय बर्बाद करते हैं। अन्य तरीके उन्हें जल्दी रुकने की कोशिश करते हैं, लेकिन वे अक्सर आलसी हो जाते हैं और गलतियाँ करने लगते हैं, या उनके लिए विशेषज्ञ को फिर से स्कूल (री-ट्रेन) जाना पड़ता है ताकि वे सीख सकें कि कब रुकना है, जो कि बहुत महंगा होता है।

समाधान: RAEE (एक "स्मार्ट लाइब्रेरियन" सिस्टम)
यह पेपर RAEE पेश करता है, जो एक सुपर-स्मार्ट लाइब्रेरियन की तरह काम करता है जो विशेषज्ञ को यह तय करने में मदद करता है कि ठीक कब रुकना है।

यह कैसे काम करता है, इसे सरल उपमाओं (analogies) में नीचे दिया गया है:

1. "समान कहानी" का तरीका (रिट्रीवल/Retrieval)

कल्पना कीजिए कि आप एक रहस्यमयी उपन्यास पढ़ रहे हैं। आपको एक सुराग मिलता है जो बहुत परिचित लगता है।

पुराना तरीका: आप सुनिश्चित होने के लिए पूरी किताब पढ़ते रहते हैं।
RAEE का तरीका: आप अपने लाइब्रेरियन को चिल्लाकर कहते हैं, "हे! मैंने पहले भी ऐसा ही सुराग देखा है!"
लाइब्रेरियन तुरंत पिछले मामलों (एक डेटाबेस) का ढेर निकालता है जहाँ ऐसे ही सुराग मिले थे। लाइब्रेरियन उन पिछले मामलों की जाँच करता है और कहता है, "9 में से 10 समान कहानियों में, जासूस ने अध्याय 5 में ही रहस्य सुलझा लिया था। आप वहीं पढ़ना बंद कर सकते हैं!"

तकनीकी शब्दों में, RAEE वर्तमान प्रश्न को देखता है, अपने पिछले प्रशिक्षण डेटा से समान प्रश्न खोजता है, और जाँचता है कि अतीत में वे समान प्रश्न कब सफलतापूर्वक हल किए गए थे।

2. "सुधारात्मक तंत्र" (जादुई हिस्सा)

यह इस पेपर का सबसे रोमांचक हिस्सा है। आमतौर पर, अर्ली एग्जिट (जल्दी रुकने वाले) तरीके एक समझौता होते हैं: गति बढ़ाओ, लेकिन बुद्धिमत्ता कम करो।
RAEE इस स्क्रिप्ट को उलट देता है। यह एक सुरक्षा जाल (safety net) की तरह काम करता है।

परिदृश्य A (आसान प्रश्न): मॉडल शुरुआत में ही आश्वस्त हो जाता है। RAEE कहता है, "यहाँ रुक जाओ!" आप समय बचाते हैं, और उत्तर अभी भी सटीक होता है।
परिदृश्य B (कठिन प्रश्न): मॉडल अंत के करीब भ्रमित हो जाता है और गलत उत्तर देने वाला होता है। लेकिन RAEE अपने डेटाबेस को देखता है और पाता कि इस विशिष्ट प्रकार के पेचीदा प्रश्न के लिए, उत्तर वास्तव में अध्याय 10 में ही स्पष्ट था, भले ही मॉडल बाद में भ्रमित हो गया हो।
- RAEE कहता है, "रुको! किताब खत्म मत करो। अध्याय 10 पर वापस जाओ। उत्तर तो वहीं था!"

परिणाम: RAEE केवल गति ही नहीं बढ़ाता; यह वास्तव में गलतियों को ठीक करता है जो पूर्ण मॉडल (full model) ने की होतीं। यह एक दूसरे विचार की तरह है जो आपके टेस्ट सबमिट करने से पहले आपकी गलतियों को पकड़ लेता है।

3. कोई नई स्कूली शिक्षा नहीं (ट्रेनिंग-फ्री)

अधिकांश अन्य तरीकों के लिए मॉडल को यह सीखने के लिए फिर से स्कूल जाना पड़ता है कि कब रुकना है। इसमें हफ्तों लग जाते हैं और बहुत पैसा खर्च होता है।
RAEE अलग है। यह मॉडल के मस्तिष्क को बिल्कुल नहीं बदलता है। यह बस एक संदर्भ पुस्तकालय (डेटाबेस) बनाता है कि "हमने पहले इसे कब सही किया था?"

उपमा: छात्र को नए नियम सिखाने के बजाय, आप उसे पिछले परीक्षाओं की एक 'चीट शीट' दे देते हैं। इसे सेटअप करना तेज़ है और इसके लिए किसी अतिरिक्त पढ़ाई की आवश्यकता नहीं है।

निचोड़ (The Bottom Line)

RA-EE को अपने मस्तिष्क की प्रोसेसिंग पावर के लिए एक GPS की तरह समझें।

RAEE के बिना: आप हर बार गंतव्य तक जाने के लिए पूरे 100 मील का सफर तय करते हैं, भले ही आप सिर्फ पास की किराने की दुकान पर जा रहे हों।
RAEE के साथ: GPS आपके इतिहास की जाँच करता है। "हे, इस किराने की दुकान की यात्रा के लिए, आप आमतौर पर मील 2 तक रास्ता जान जाते हैं। चलिए वहीं रुक जाते हैं।" लेकिन यदि आप किसी नए शहर जा रहे हैं और रास्ता भटक जाते हैं, तो GPS समान मार्गों की जाँच करता है और कहता है, "वास्तव में, आप मील 10 पर मोड़ जानते थे, उसके आगे गाड़ी चलाना जारी न रखें!"

यह क्यों महत्वपूर्ण है:

तेज़: यह कार्यों को बहुत तेज़ी से पूरा करता है।
स्मार्ट: यह अक्सर पूर्ण गति से चलने वाले मॉडल की तुलना में बेहतर उत्तर प्राप्त करता है क्योंकि यह उस "ओवरथिंकिंग" से बचता है जो गलतियों का कारण बनती है।
सस्ता: यह बिना री-ट्रेनिंग के भारी मात्रा में बिजली और कंप्यूटिंग शक्ति बचाता है।

संक्षेप में, RAEE AI को लापरवाह हुए बिना कुशल बनना सिखाता है, अपने पिछले अनुभवों के ज्ञान का उपयोग करके यह जानना कि कब रुकना है।

यहाँ शोध पत्र "RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

लार्ज लैंग्वेज मॉडल्स (LLMs) और अन्य डीप न्यूरल नेटवर्क्स को तैनात करने में उच्च कम्प्यूटेशनल ओवरहेड और मेमोरी की आवश्यकताएं एक बड़ी बाधा हैं। अर्ली एग्जिट (Early Exit) तंत्र इसे तब संबोधित करते हैं जब एक कॉन्फिडेंस मानदंड (confidence criterion) पूरा हो जाता है, जिससे पूरे मॉडल को चलाने के बजाय मध्यवर्ती परतों (intermediate layers) पर ही इन्फरेंस को गतिशील रूप से समाप्त किया जा सकता है।

हालाँकि, मौजूदा अर्ली एग्जिट फ्रेमवर्क महत्वपूर्ण सीमाओं से ग्रस्त हैं:

ट्रेनिंग-आधारित विधियाँ (Training-based methods): इनके लिए आंतरिक क्लासिफायर और बैकबोन के संयुक्त अनुकूलन (joint optimization) की आवश्यकता होती है, जिससे भारी ट्रेनिंग ओवरहेड होता है।
सेमी-ट्रेनिंग विधियाँ (Semi-training methods): ये बैकबोन को फ्रीज कर देती हैं लेकिन मैनुअल फीचर इंजीनियरिंग और हल्के क्लासिफायर पर निर्भर करती हैं, जिनमें अक्सर सामान्यीकरण (generalization) की कमी होती है।
ट्रेनिंग-फ्री विधियाँ (Training-free methods): ये ह्यूरिस्टिक मानदंडों (जैसे एंट्रॉपी थ्रेसहोल्ड) का उपयोग करती हैं जो अनुकूलन क्षमता में कमी रखते हैं, जिससे अक्सर पूर्ण मॉडल की तुलना में प्रदर्शन में गिरावट (accuracy loss) आती है।

अधिकांश वर्तमान दृष्टिकोण एक ऐसा समझौता (trade-off) करने के लिए मजबूर करते हैं जहाँ गति प्राप्त करने के लिए सटीकता (accuracy) का त्याग करना पड़ता है। यह शोध पत्र इस समझौते को तोड़ने के लिए एक मजबूत, ट्रेनिंग-फ्री अर्ली एग्जिट फ्रेमवर्क विकसित करने का लक्ष्य रखता है जो मॉडल की सटीकता को बनाए रखते हुए या उसमें सुधार करते हुए इन्फरेंस को तेज करता है।

2. मुख्य अवलोकन और प्रेरणा (Key Observations & Motivation)

लेखक दो महत्वपूर्ण अवलोकन प्रस्तावित करते हैं जो पारंपरिक अर्ली एग्जिट प्रतिमानों को चुनौती देते हैं:

सुधारात्मक तंत्र के रूप में अर्ली एग्जिट (Early Exit as a Corrective Mechanism): मध्यवर्ती परतें कभी-कभी सही भविष्यवाणी कर सकती हैं भले ही अंतिम परत विफल हो जाए। यदि मॉडल एक ऐसी मध्यवर्ती परत पर जल्दी बाहर निकल जाता है जहाँ भविष्यवाणी सही है, तो यह पूर्ण मॉडल के अंतिम आउटपुट को "सुधार" (correct) सकता है।
समान डेटा की निरंतरता (Consistency of Similar Data): अर्थ संबंधी रूप से समान इनपुट अत्यधिक सुसंगत एग्जिट व्यवहार प्रदर्शित करते हैं। यदि एक विशिष्ट इनपुट $x$ को परत $L$ पर सबसे अच्छा बाहर निकलना चाहिए, तो एम्बेडिंग स्पेस में उसके निकटतम पड़ोसी भी उच्च विश्वास के साथ परत $L$ (या एक समान परत) पर ही सबसे अच्छा बाहर निकलने की संभावना रखते हैं।

3. कार्यप्रणाली: RAEE फ्रेमवर्क (Methodology: RAEE Framework)

RAEE (Retrieval-Augmented Early Exit) एक ट्रेनिंग-फ्री फ्रेमवर्क है जो इष्टतम एग्जिट लेयर निर्धारित करने के लिए एक बाहरी रिट्रीवल डेटाबेस का लाभ उठाता है। यह अर्ली एग्जिट समस्या को एक डिस्ट्रीब्यूशन प्रेडिक्शन समस्या (distribution prediction problem) के रूप में मानता है।

A. रिट्रीवल डेटाबेस का निर्माण (Offline Phase)

डेटा संग्रह (Data Collection): फ्रेमवर्क एक ट्रेनिंग डेटासेट $D$ को बैकबोन मॉडल $M$ के माध्यम से प्रोसेस करता है।
फीचर एक्सट्रैक्शन (Feature Extraction): प्रत्येक इनपुट $x_i$ के लिए, मॉडल एम्बेडिंग्स (एक बाहरी एनकोडर या बैकबोन की एम्बेडिंग लेयर का उपयोग करके) कंप्यूट करता है जो Keys के रूप में कार्य करते हैं।
वैल्यू कंस्ट्रक्शन (Value Construction): मॉडल सभी परतों के माध्यम से जाता है। प्रत्येक परत $j$ के लिए जहाँ मध्यवर्ती भविष्यवाणी $\hat{y}_j$ ग्राउंड ट्रुथ $y_i$ से मेल खाती है, जोड़ी $(j, p_j)$ (लेयर इंडेक्स और प्रेडिक्शन प्रोबेबिलिटी) को रिकॉर्ड किया जाता है।
स्टोरेज (Storage): ये जोड़ियाँ इन इनपुट कीज़ से जुड़ी Values बनाती हैं। इन की-वैल्यू पेयर्स को स्टोर करने के लिए एक एप्रोक्सिमेट नियरएस्ट नेबर (ANN) इंडेक्स (जैसे FAISS) बनाया जाता है।
- नोट: इस प्रक्रिया के लिए बैकबोन मॉडल के पैरामीटर्स को अपडेट करने की आवश्यकता नहीं है।

B. इन्फरेंस चरण (Online Phase)

क्वेरी जनरेशन (Query Generation): एक नए इनपुट $x$ के लिए, फ्रेमवर्क एक एम्बेडिंग क्वेरी जेनरेट करता है।
रिट्रीवल (Retrieval): सिस्टम डेटाबेस से शीर्ष- $k$ निकटतम पड़ोसियों को पुनः प्राप्त करता है।
डिस्ट्रीब्यूशन एप्रोक्सिमेशन (Distribution Approximation): RAEE पड़ोसियों से एग्जिट जानकारी को एकत्रित करता है। यह $k$ पड़ोसियों के माध्यम से लेयर $l$ पर एग्जिट करने की संभाव्यता वितरण $P(z=l|x)$ का अनुमान लगाता है, जो पड़ोसियों के वैध एग्जिट लेयर्स का एक भारित योग (weighted sum) है, जहाँ भार क्वेरी और पड़ोसी के बीच की दूरी के व्युत्क्रमानुपाती (inversely proportional) होता है।
$P(z = l | x) = \sum_{i=1}^{k} P(v_i | x) \cdot S_i$
जहाँ $S_i$ पड़ोसी $i$ के योगदान को गिनता है यदि उसके पास थ्रेसहोल्ड $\tau$ से ऊपर की संभावना के साथ लेयर $l$ पर एक वैध एग्जिट है।
एग्जिट निर्णय (Exit Decision): फ्रेमवर्क उस लेयर $l$ का चयन करता है जो इस संभाव्यता को अधिकतम करती है: $f(x) = \arg\max_l P(z=l|x)$ । यदि कई परतों में समान अधिकतम संभाव्यता है, तो सबसे प्रारंभिक परत चुनी जाती है।
निष्पादन (Execution): मॉडल केवल चयनित लेयर $l$ तक ही फॉरवर्ड रन करता है, और मध्यवर्ती आउटपुट को अंतिम प्रेडिक्शन हेड (जैसे LM Head) को पास कर दिया जाता है।

4. मुख्य योगदान (Key Contributions)

नया स्वरूप (Novel Formulation): यह पेपर अर्ली एग्जिट को एक डिस्ट्रीब्यूशन प्रेडिक्शन समस्या के रूप में मॉडल करता है, यह प्रदर्शित करता है कि किसी क्वेरी का एग्जिट डिस्ट्रीब्यूशन उसके अर्थ संबंधी रूप से समान पड़ोसियों के एग्जिट व्यवहार द्वारा प्रभावी रूप से अनुमानित किया जा सकता है।
RAEE फ्रेमवर्क: एक मजबूत, ट्रेनिंग-फ्री रिट्रीवल-ऑगमेंटेड फ्रेमवर्क प्रस्तावित करता है। यह आंतरिक क्लासिफायर को प्रशिक्षित करने या बैकबोन को फाइन-ट्यून करने की आवश्यकता को समाप्त करता है, जिससे डिप्लॉयमेंट ओवरहेड काफी कम हो जाता है।
दोहरा लाभ (स्पीड + सटीकता): पारंपरिक तरीकों के विपरीत जो गति के लिए सटीकता का त्याग करते हैं, RAEE एक सुधारात्मक तंत्र (corrective mechanism) के रूप में कार्य करता है। उन मामलों से एग्जिट जानकारी का लाभ उठाकर जहाँ मध्यवर्ती परतें सफल रहीं लेकिन पूर्ण मॉडल विफल हो सकता है (या इसके विपरीत), यह समग्र सटीकता में सुधार करता है और साथ ही इन्फरेंस को तेज करता है।

5. प्रयोगात्मक परिणाम (Experimental Results)

लेखकों ने विभिन्न बैकबोन मॉडल्स: RoBERTa-Large, ElasticBERT-Large, T5-Large, Llama-3-8B, और Gemma-7B का उपयोग करके आठ डाउनस्ट्रीम कार्यों (GLUE बेंचमार्क) पर RAEE का मूल्यांकन किया।

सटीकता में सुधार (Accuracy Improvement):
- RAEE ने लगातार अत्याधुनिक अर्ली एग्जिट विधियों (जैसे DeeBERT, AdaInfer, HashEE, CALM) को पछाड़ दिया।
- RoBERTa-Large के साथ, RAEE ने 63.41 की औसत सटीकता प्राप्त की, जो बेसलाइन तुलना विधियों (~42.99 औसत) से काफी बेहतर है।
- विशेष रूप से, RAEE ने कभी-कभी पूर्ण मॉडल की सटीकता को भी पार कर लिया, जो "सुधारात्मक तंत्र" पर इसके परिकल्पना (hypothesis) की पुष्टि करता है।
इन्फरेंस दक्षता (Inference Efficiency):
- बिलियन-पैरामीटर मॉडल्स (Llama-3-8B, Gemma-7B) के लिए, RAEE ने पूर्ण मॉडल चलाने की तुलना में इन्फरेंस लेटेंसी को लगभग 50% कम कर दिया, जबकि सटीकता को बनाए रखा या सुधारा।
- छोटे मॉडल्स के लिए, इसने अन्य अर्ली एग्जिट विधियों के समान दक्षता हासिल की लेकिन बेहतर सटीकता प्रदान की।
एब्लेशन स्टडीज (Ablation Studies):
- रिट्रीवल साइज ( $k$ ): प्रदर्शन $k$ के 12 तक बढ़ने के साथ सुधरा, जिसके बाद कम प्रासंगिक पड़ोसियों के शोर (noise) के कारण मामूली गिरावट आई।
- डेटाबेस साइज: डेटाबेस के लिए उपयोग किए गए बड़े ट्रेनिंग सेट्स ने बेहतर सामान्यीकरण और सटीकता प्रदान की।
- आउट-ऑफ-डोमेन (OOD): जब इसे (WikiText से निर्मित डेटाबेस का उपयोग करके) समराइजेशन कार्यों (CNN/DailyMail, XSum) पर टेस्ट किया गया, तब भी RAE ने प्रदर्शन में सुधार किया और लेयर के उपयोग को कम किया, जो इसकी मजबूती (robustness) को दर्शाता है।

6. महत्व (Significance)

RAEE, LLMs के कुशल इन्फरेंस के लिए एक प्रतिमान परिवर्तन (paradigm shift) का प्रतिनिधित्व करता है।

ट्रेनिंग ओवरहेड को समाप्त करता है: यह आंतरिक क्लासिफायर को प्रशिक्षित करने की बाधा को दूर करता है, जिससे इसे किसी भी प्री-ट्रेंड मॉडल पर तुरंत तैनात किया जा सकता है।
सटीकता-गति के समझौते को तोड़ता है: यह सिद्ध करता है कि अर्ली एग्जिट का अर्थ प्रदर्शन में गिरावट नहीं है; सही रिट्रीवल रणनीति के साथ, यह गहरी परतों द्वारा की गई त्रुटियों को सुधारकर प्रदर्शन को बढ़ा सकता है।
स्केलेबिलिटी: यह विधि उन बड़े मॉडल्स (7B-8B पैरामीटर्स) के लिए अच्छी तरह से स्केल करती है जहाँ इन्फरेंस लागत सबसे अधिक होती है, जो LLMs के वास्तविक दुनिया के डिप्लॉयमेंट के लिए एक व्यावहारिक समाधान प्रदान करती है।

संक्षेप में, RAEE "समान डेटा की बुद्धिमत्ता" का लाभ उठाकर गतिशील और बुद्धिमानी से इन्फरेंस पथ को छोटा करता है, जिससे उच्च गति, उच्च सटीकता और शून्य प्रशिक्षण लागत का एक दुर्लभ संयोजन प्राप्त होता है।