🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

यह शोध पत्र टोकन-सिलेक्टिव अटेंशन (TSA) को प्रस्तुत करता है, जो एक हल्का, एंड-टू-एंड डिफरेंशिएबल गेटिंग मैकेनिज्म है जो ट्रांसफॉर्मर्स को संदर्भ-आधारित कठिनाई के आधार पर व्यक्तिगत टोकन के लिए अनावश्यक लेयर गणनाओं को गतिशील रूप से छोड़ने में सक्षम बनाता है, जिससे बिना किसी स्पष्ट डेप्थ रेगुलराइजेशन की आवश्यकता के और न्यूनतम गुणवत्ता हानि के साथ महत्वपूर्ण दक्षता लाभ प्राप्त होता है।

मूल लेखक: Ahmed Abdelmuniem Abdalla Mohammed

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

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

CC BY 4.0

मूल लेखक: Ahmed Abdelmuniem Abdalla Mohammed

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

कल्पना कीजिए कि आप एक फैक्ट्री चला रहे हैं जहाँ असेंबली लाइन से आने वाले हर एक उत्पाद को बिल्कुल एक जैसा उपचार मिलता है। चाहे वह एक साधारण, उत्तम विजेट हो या एक जटिल, टूटा हुआ विजेट, हर आइटम हर एक स्टेशन पर ठीक उतना ही समय बिताता है। यह तरीका वर्तमान में मानक AI मॉडल (ट्रांसफॉर्मर्स) के काम करने का है: वे वाक्य के हर शब्द को समान संख्या में "सोचने वाली परतों" (thinking layers) के माध्यम से प्रोसेस करते हैं, चाहे वह शब्द समझने में कितना भी आसान या कठिन क्यों न हो।

यह पेपर एक नई प्रणाली पेश करता है जिसे टोकन-सिलेक्टिव अटेंशन (TSA) कहा जाता है। TSA को एक फैक्ट्री में हर स्टेशन पर लगे एक स्मार्ट, स्वचालित गेटकीपर (gatekeeper) के रूप में समझें।

समस्या: "एक ही आकार सबके लिए" वाली फैक्ट्री

एक मानक AI मॉडल में, यदि आप उससे "To be, or not to be" लिखने के लिए कहते हैं, तो मॉडल "To" शब्द और "question" शब्द दोनों के साथ समान मात्रा में दिमागी शक्ति का उपयोग करता है।

आसान शब्द (जैसे "the" या "is") सरल विजेट की तरह हैं। उन्हें बहुत अधिक प्रोसेसिंग की आवश्यकता नहीं होती।
कठिन शब्द (जैसे जटिल अवधारणाएं या दुर्लभ नाम) टूटे हुए विजेट की तरह हैं। उन्हें ठीक करने के लिए अतिरिक्त समय और ध्यान की आवश्यकता होती है।

वर्तमान में, फैक्ट्री आसान विजेट्स को भी उतनी ही तीव्रता से प्रोसेस करने में ऊर्जा बर्बाद करती है जितनी कि कठिन विजेट्स को।

समाधान: स्मार्ट गेटकीपर (TSA)

लेखकों ने AI के प्रत्येक लेयर के बीच एक छोटा, हल्का "गेटकीपर" (एक छोटा न्यूरल नेटवर्क) जोड़ा है। यह गेटकीपर प्रत्येक शब्द (टोकन) को व्यक्तिगत रूप से देखता है और पूछता है: "क्या इस शब्द को अगले प्रोसेसिंग स्टेशन से गुजरने की आवश्यकता है, या यह इसे छोड़ सकता है?"

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

व्यवहार में यह कैसे काम करता है

शोधकर्ताओं ने इसका परीक्षण दो मुख्य चीजों पर किया:

सरल तर्क पहेलियाँ (Logic Puzzles): जब AI को संख्याओं की एक सूची को कॉपी करने के लिए कहा गया, तो गेटकीपर ने महसूस किया, "यह आसान है," और उसने लगभग 65% काम छोड़ दिया। जब उसे संख्याओं को सॉर्ट (क्रम में लगाना) करने के लिए कहा गया (जो कि कठिन है), तो उसने केवल 27% काम छोड़ा। उसने स्वाभाविक रूप से समझ लिया कि कठिन कार्यों के लिए अधिक चरणों की आवश्यकता होती है।
कहानियां लिखना: जब AI को शेक्सपियर की तरह लिखने या विकिपीडिया लेखों का सारांश लिखने के लिए कहा गया, तो AI ने अपनी कंप्यूटिंग शक्ति में 14% से 23% तक की बचत की।
- इसने सीखा कि विराम चिह्न, स्पेस और सामान्य शब्दों को आसानी से प्रोसेस किया जा सकता है।
- इसने सीखा कि अद्वितीय सामग्री वाले शब्दों को पूरे "फैक्ट्री" उपचार की आवश्यकता होती है।

परिणाम: तेज़ और स्मार्ट

गुणवत्ता में कोई कमी नहीं: AI ने मानक मॉडल की तरह ही अच्छा लिखा, लेकिन इसने काफी कम ऊर्जा (कंप्यूटिंग पावर) का उपयोग किया।
"अर्ली एग्जिट" (Early Exit) से बेहतर: अन्य तरीके कोशिश करते हैं कि यदि AI "आत्मविश्वास" महसूस करे तो पूरी प्रक्रिया को ही जल्दी रोक दें। TSA अधिक स्मार्ट है; यह व्यक्तिगत शब्दों को अनावश्यक चरणों से गुजरने से रोकता है जबकि कठिन शब्दों को आगे बढ़ने देता है। पेपर में पाया गया कि जब TSA को समान ऊर्जा बचाने के लिए सेट किया गया था, तो इसने पुरानी विधियों की तुलना में बेहतर परिणाम दिए।
वास्तविक गति: जब शोधकर्ताओं ने वास्तव में कंप्यूटर पर कोड चलाया, तो उन्होंने वास्तविक स्पीडअप (लगभग 2.3% तेज़) देखा क्योंकि कंप्यूटर को वास्तव में उन शब्दों के लिए गणित नहीं करना पड़ा जिन्हें छोड़ दिया गया था।

निचोड़

यह पेपर एक तरीका प्रस्तुत करता है जिससे AI मॉडल को स्मार्ट तरीके से "आलसी" बनाया जा सके। हर शब्द को समान मात्रा में काम करने के लिए मजबूर करने के बजाय, TSA यह तय करने की अनुमति देता है कि कौन से शब्द आसान हैं और जिन्हें अतिरिक्त मदद की आवश्यकता है। यह एक ऐसी फैक्ट्री की तरह है जहाँ आसान उत्पाद लाइन से तेजी से निकल जाते हैं, जबकि कठिन उत्पादों को पूरा वीआईपी (VIP) उपचार मिलता है, जिससे गुणवत्ता से समझौता किए बिना समय और ऊर्जा की बचत होती है।

तकनीकी सारांश: ट्रांसफॉर्मर्स में सीखे गए टोकन रूटिंग के माध्यम से अनुकूलित गणना गहराई (Adaptive Computation Depth)

समस्या विवरण

मानक ट्रांसफॉर्मर आर्किटेक्चर अनुक्रम के प्रत्येक टोकन पर, उस टोकन की प्रासंगिक कठिनाई या सूचना घनत्व की परवाह किए बिना, परतों (layers) की एक निश्चित संख्या लागू करते हैं। यह डिज़ाइन प्रति-टोकन अनुकूलन क्षमता के बजाय वास्तुशिल्प सरलता को प्राथमिकता देता है, जिसके परिणामस्वरूप महत्वपूर्ण गणना अक्षमता उत्पन्न होती है। व्यवहार में, अनुमानित संदर्भों में सामान्य टोकन को नए निर्माणों में दुर्लभ टोकन की तुलना में बहुत कम प्रोसेसिंग की आवश्यकता होती है, फिर भी दोनों को प्रत्येक फॉरवर्ड पास के दौरान समान कंप्यूट संसाधन प्राप्त होते हैं। मौजूदा समाधान, जैसे कि RNNs के लिए एडैप्टिव कंप्यूटेशन टाइम (ACT), यूनिवर्सल ट्रांसफॉर्मर्स, मिक्सचर-ऑफ-डेप्थ (MoD), और अर्ली एग्जिट (Early Exit) तंत्र, इस समस्या को संबोधित करने का प्रयास करते हैं, लेकिन वे अक्सर हार्ड सिलेक्शन, संचित बजट, या पोस्ट-हॉक कॉन्फिडेंस थ्रेशोल्ड पर निर्भर करते हैं जो पूरी तरह से डिफरेंशिएबल या इष्टतम नहीं हो सकते हैं।

कार्यप्रणाली: टोकन-सिलेक्टिव अटेंशन (TSA)

यह शोध पत्र टोकन-सिलेक्टिव अटेंशन (TSA) का प्रस्ताव करता है, जो एक तंत्र है जो क्रमिक ट्रांसफॉर्मर ब्लॉक्स के बीच रेसिडुअल अपडेट पर एक सीखा गया, प्रति-टोकन गेट पेश करता है।

आर्किटेक्चर और रूटिंग

TSA एक प्री-नॉर्म डिकोडर-ओनली ट्रांसफॉर्मर में प्रत्येक ब्लॉक $f_l$ (जहाँ $l = 0, \dots, L-2$ ) के बाद एक हल्का राउटर $r_l$ सम्मिलित करता है। स्टेम ब्लॉक $f_0$ हमेशा बिना शर्त निष्पादित होता है।

राउटर डिज़ाइन: प्रत्येक राउटर एक दो-परत वाला मल्टी-लेयर पर्सपेक्ट्रॉन (MLP) है जिसका आउटपुट सिग्मॉइड (sigmoid) होता है, जो प्रत्येक टोकन के लिए एक निरंतर हॉल्टिंग प्रोबेबिलिटी (halting probability) $p_l \in (0, 1)$ उत्पन्न करता है। इसका हिडन डायमेंशन $d/4$ (16 पर फ्लोर किया गया) है, जो बेस मॉडल में लगभग 1.7% पैरामीटर ओवरहेड जोड़ता है।
गेटेड अपडेट: आगामी ब्लॉक $f_{l+1}$ $f_{l + 1}$ के लिए रेसिडुअल अपडेट को $(1 - p_l)$ $(1 - p_{l})$ द्वारा स्केल किया जाता है।
- यदि $p_l = 0$ है, तो अपडेट मानक होता है।
- यदि $p_l = 1$ है, तो स्टेट अपरिवर्तित रहती है, जो प्रभावी रूप से ब्लॉक को स्किप कर देती है।
- इंटरपोलेशन सुचारू (smooth) है, जो यह सुनिश्चित करता है कि तंत्र पूरी तरह से डिफरेंशिएबल है और इसके लिए स्ट्रैथ-थ्रू एस्टिमेटर्स, गम्बल सैंपलिंग, या रिइन्फोर्समेंट लर्निंग की आवश्यकता नहीं है।

प्रशिक्षण और रेगुलाइजेशन

प्रशिक्षण के दौरान, सभी परतें पूरी तरह से निष्पादित होती हैं; गेट रेसिडुअल अपडेट को स्केल करता है लेकिन गणना को स्किप नहीं करता है। कुल लॉस $L = L_{task} + L_{depth}$ है।

डेप्थ रेगुलाइजेशन: एक टर्म $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ जोड़ा गया है जो धीरे से अर्ली हॉल्टिंग को प्रोत्साहित करता है।
इमर्जेंट बिहेवियर (Emergent Behavior): विशेष रूप से, शोध पत्र यह प्रदर्शित करता है कि रूटिंग तंत्र तब भी ऑपरेशन को स्किप करना सीख जाता है जब $\lambda = 0$ (कोई स्पष्ट डेप्थ प्रेशर नहीं) हो, जो केवल टास्क-लॉस ग्रेडिएंट द्वारा संचालित होता है।

इन्फरेंस (Inference)

इन्फरेंस के समय, तंत्र "स्पार्स" (sparse) मोड में स्विच हो जाता है। जिन टोकन्स के लिए $p_l > 0.5$ है, उन्हें गेदर/स्कैटर ऑपरेशंस के माध्यम से फीड-फॉरवर्ड नेटवर्क (FFN) से बाहर रखा जाता है, जबकि सटीक की-वैल्यू सिमेंटिक्स को बनाए रखने के लिए अटेंशन डेंस रहता है। इससे वास्तविक FLOPs बचत और वॉल-क्लॉक स्पीडअप प्राप्त होता है।

मुख्य योगदान

डिफरेंशिएबल टोकन रूटिंग: रेसिडुअल अपडेट्स पर एक सरल, निरंतर सॉफ्ट गेट जो बेस ट्रांसफॉर्मर में संरचनात्मक परिवर्तन के बिना प्रति-टोकन, प्रति-लेयर रूटिंग की अनुमति देता है।
ग्रेडिएंट-ड्रिवन इमर्जेंस: प्रमाण कि टास्क-लॉस ग्रेडिएंट से ही कठिनाई-अनुपातिक रूटिंग स्वाभाविक रूप से उभरती है। बिना डेप्थ रेगुलाइजेशन ( $\lambda=0$ ) के भी, राउटर लगभग 20% टोकन-लेयर ऑपरेशंस को स्किप करना सीख जाता है।
एफिशिएंसी-क्वालिटी ट्रेड-ऑफ: कैरेक्टर-लेवल लैंग्वेज मॉडलिंग पर वैलिडेशन कि 14–23% टोकन-लेयर ऑपरेशंस (TLOps) की बचत के साथ 0.5% से कम क्वालिटी लॉस देखा गया।
अर्ली एग्जिट से श्रेष्ठता: समान दक्षता पर, TSA अर्ली एग्जिट बेसलाइन्स की तुलना में कम वैलिडेशन लॉस प्राप्त करता है, जो यह सुझाव देता है कि एंड-टू-एंड सीखे गए रूटिंग निर्णय पोस्ट-हॉक कॉन्फिडेंस थ्रेशोल्डिंग की तुलना में बेहतर हैं।
रियल-वर्ल्ड स्पीडअप: स्पार्स इन्फरेंस के माध्यम से कमोडिटी हार्डवेयर (Apple M1 Pro) पर वास्तविक वॉल-क्लॉक स्पीडअप (2.3% तक) का प्रदर्शन।

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

सिंथेटिक टास्क

कॉपी और सॉर्ट टास्क (लंबाई-10 अनुक्रम) पर:

कॉपी टास्क: राउटर ने सीखा कि कार्य एक आइडेंटिटी मैपिंग है, जिससे बिना किसी सटीकता हानि के 54.9% TLOps को स्किप किया गया ( $\alpha = 0.341$ )।
सॉर्ट टास्क: राउटर ने अधिक कंप्यूट आवंटित किया, 22.5% TLOps को स्किप किया ( $\alpha = 0.730$ ), जो उच्च कठिनाई को दर्शाता है।

कैरेक्टर-लेवल लैंग्वेज मॉडलिंग

Tiny-Shakespeare: TSA ने 22.8% TLOps ( $\alpha = 0.726$ ) की बचत की, जिसमें वैलिडेशन लॉस में केवल 0.006 नैट्स (+0.4%) की वृद्धि हुई।
enwik8: विविध कॉर्पस (Wikipedia) पर, TSA ने बिना किसी गुणवत्ता गिरावट के 13.9% TLOps ( $\alpha = 0.833$ ) की बचत की (मामूली बेहतर वैलिडेशन लॉस)।
रोबस्टनेस (Robustness): तंत्र दो ऑर्डर्स ऑफ मैग्नीट्यूड ( $[0, 0.1]$ ) तक $\lambda$ के लिए रोबस्ट है। $\lambda=0.05$ पर, इसने नगण्य क्वालिटी कॉस्ट के साथ 50.4% TLOps बचत हासिल की।

अर्ली एग्जिट के साथ तुलना

एक मैच किए गए एक्टिव फ्रैक्शन ( $\alpha \approx 0.726$ ) पर, TSA ने अर्ली एग्जिट बेसलाइन की तुलना में 0.71% कम वैलिडेशन लॉस प्राप्त किया। यह इंगित करता है कि एंड-टू-एंड सीखे गए रूटिंग निर्णय, इन्फरेंस के समय कॉन्फिडेंस थ्रेशोल्ड लागू करने की तुलना में बेहतर निर्णय देते हैं।

वॉल-क्लॉक थ्रूपुट

Soft-TSA (Training): $\alpha$ के बावजूद ~1% ओवरहेड जोड़ा।
Sparse-TSA (Inference): Tiny-Shakespeare ( $\alpha=0.726$ ) पर 2.3% की स्पीडअप और enwik8 ( $\alpha=0.833$ ) पर ब्रेक-इवन प्रदर्शन प्राप्त किया। CPU-GPU सिंक्रोनाइज़ेशन ओवरहेड को पार करने के लिए बैच साइज $\ge 64$ की आवश्यकता थी।

महत्व और दावे

शोध पत्र का दावा है कि TSA, फिक्स्ड-डेप्थ ट्रांसफॉर्मर्स की अक्षमता का एक न्यूनतम, वास्तुशिल्प रूप से सरल समाधान प्रदान करता है। इसका प्राथमिक महत्व यह प्रदर्शित करना है कि अनुकूलित गणना गहराई को स्पष्ट पर्यवेक्षण या जटिल रिइन्फोर्समेंट लर्निंग संकेतों के बिना सीधे टास्क ऑब्जेक्टिव से सीखा जा सकता है।

लेखक इस बात पर जोर देते हैं कि राउटर एक "लर्नड नॉइज़ गेट" के रूप में कार्य करता है, जो उन अपडेट्स को कम करता है जब रेसिडुअल योगदान रेडंडेंट (redundant) होते हैं। यह मॉडल को टोकन की कठिनाई के आधार पर गतिशील रूप से कंप्यूट आवंटित करने की अनुमति देता है, जिससे न्यूनतम क्वालिटी ट्रेड-ऑफ के साथ महत्वपूर्ण दक्षता लाभ (सिंथेटिक टास्क पर 55% तक, लैंग्वेज मॉडलिंग पर 14-23%) प्राप्त होता है। यह कार्य बताता है कि एंड-टू-एंड डिफरेंशिएबल रूटिंग, डिस्क्रीट सिलेक्शन या पोस्ट-हॉक थ्रेशोल्डिंग की तुलना में एडैप्टिव डेप्थ के लिए एक अधिक प्रभावी रणनीति है, जो कमोडिटी हार्डवेयर पर अधिक कुशल इन्फरेंस का मार्ग प्रशस्त करती है।

शोध पत्र अपनी स्केल के संबंध में विनम्र है, यह नोट करते हुए कि सभी प्रयोग लगभग 5-6 मिलियन पैरामीटर्स वाले मॉडल्स पर किए गए थे, और बड़े मॉडल्स (10M-100M+) तक स्केलिंग व्यवहार अभी भी एक खुला प्रश्न है। इसके अतिरिक्त, वर्तमान स्पीडअप बैच-साइज पर निर्भर है, और अटेंशन स्पर्सिटी (attention sparsity) को अभी तक लागू नहीं किया गया है।