🤖 machine learning

Stochastic Attention: Connectome-Inspired Randomized Routing for Expressive Linear-Time Attention

फ्रूट फ्लाई के कनेक्टोम से प्रेरित होकर, यह शोध पत्र स्टोकेस्टिक अटेंशन (Stochastic Attention) का प्रस्ताव करता है, जो एक रैंडमाइज्ड रूटिंग तंत्र है जो स्लाइडिंग-विंडो अटेंशन को बेहतर बनाने के लिए लॉगैरिद्मिक डेप्थ में ग्लोबल रिसेप्टिव फील्ड्स प्राप्त करता है और लीनियर-टाइम कॉम्प्लेक्सिटी बनाए रखता है, जो प्री-ट्रेनिंग और ट्रेनिंग-फ्री इन्फरेंस परिदृश्यों दोनों में उत्कृष्ट प्रदर्शन प्रदर्शित करता है।

मूल लेखक: Zehao Jin, Yanan Sui

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

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

CC BY 4.0

मूल लेखक: Zehao Jin, Yanan Sui

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

मुख्य विचार: AI के लिए एक "रैंडम शफल" (Random Shuffle) ट्रिक

कल्पना कीजिए कि आप एक बहुत लंबी किताब पढ़ने की कोशिश कर रहे हैं, लेकिन आपको एक बार में केवल 10 शब्दों की एक छोटी खिड़की (window) देखने की अनुमति है। कई वर्तमान AI मॉडल ऊर्जा बचाने के लिए इसी तरह काम करते हैं; वे स्लाइडिंग विंडो अटेंशन (Sliding Window Attention - SWA) का उपयोग करते हैं। वे पास-पास वाले शब्दों को देखते हैं, जो स्थानीय वाक्यों को समझने के लिए तो बेहतरीन है, लेकिन यह याद रखने के मामले में बहुत खराब है कि 50 पन्ने पहले क्या हुआ था।

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

यह शोध पत्र एक नई विधि प्रस्तावित करता है जिसे स्टोकेस्टिक अटेंशन (Stochastic Attention - SA) कहा जाता है। यह उस लंबी किताब को लेने, उसके पन्नों को यादृच्छिक रूप से शफल (shuffle) करने, आपके 10-शब्दों के विंडो को पढ़ने और फिर पन्नों को वापस उनके सही क्रम में रखने जैसा है।

यह कैसे काम करता है: "शफल और रीड" (Shuffle and Read) सादृश्य

1. "लोकल विंडो" की समस्या
एक मानक AI मॉडल को एक ऐसे व्यक्ति के रूप में सोचें जो आंखों पर पट्टी बांधकर किताब पढ़ रहा है, जिससे वह केवल वर्तमान शब्द और उससे पहले के 10 शब्दों को ही देख सकता है।

समस्या: यदि उस व्यक्ति को पेज 1 के किसी शब्द को पेज 100 के किसी शब्द से जोड़ने की आवश्यकता है, तो वह ऐसा नहीं कर सकता। उसे पूरी किताब को पेज-दर-पेज पढ़ना होगा, जिसमें बहुत समय लगता है।

2. फल मक्खी (Fruit Fly) से प्रेरणा
फल मक्खी का मस्तिष्क दक्षता का एक उत्कृष्ट नमूना है। इसमें लगभग 1,30,000 न्यूरॉन्स होते हैं। अधिकांश केवल अपने निकटतम पड़ोसियों से जुड़े होते हैं (लोकल क्लस्टरिंग)। लेकिन, कुछ यादृच्छिक कनेक्शन (random connections) पूरे मस्तिष्क में लंबी दूरी तय करते हैं (long-range shortcuts)।

परिणाम: भले ही मस्तिष्क ज्यादातर स्थानीय (local) हो, एक सिग्नल केवल कुछ ही "हॉप्स" (लग्स/कदमों) में एक छोर से दूसरे छोर तक जा सकता है (लगभग 4 कदम)। यह एक "स्मॉल वर्ल्ड" नेटवर्क है।

3. समाधान: स्टोकेस्टिक अटेंशन (Stochastic Attention - SA)
लेखकों ने महसूस किया कि वे एक जटिल नया मस्तिष्क बनाए बिना फल मक्खी के इन शॉर्टकट्स की नकल कर सकते हैं। वे इसे तीन-चरणीय ट्रिक के साथ करते हैं:

चरण 1: शफल (Shuffle)। AI के अपने शब्दों के "विंडो" को देखने से पहले, यह पूरे वाक्य के क्रम को यादृच्छिक रूप से बिखेर (scramble) देता है।
चरण 2: रीड (Read)। AI अपने 10 शब्दों के छोटे विंडो को देखता है। चूंकि वाक्य को शफल किया गया था, इसलिए वे 10 शब्द मूल कहानी के शुरुआत, मध्य और अंत से आ सकते हैं। अब AI दूर के हिस्सों को ऐसे "देख" रहा है जैसे कि वे पड़ोसी हों।
चरण 3: अनशफल (Unshuffle)। AI द्वारा जानकारी को प्रोसेस करने के बाद, वह शब्दों को उनके मूल क्रम में वापस रख देता है ताकि वाक्य का अर्थ बना रहे।

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

"दोनों दुनियाओं का सर्वश्रेष्ठ" संयोजन (The "Best of Both Worlds" Combo)

शोध पत्र यह भी सुझाव देता है कि केवल सब कुछ शफल करना पर्याप्त नहीं है। आपको शब्दों के स्थानीय प्रवाह (व्याकरण, वाक्य संरचना) को भी समझना होगा। इसलिए, उन्होंने एक गेटेड SA + SWA सिस्टम बनाया है।

इसे एक दो-लेन वाले हाईवे के रूप में सोचें:

लेन 1 (SWA): मानक लेन जहाँ कारें (शब्द) अपने स्थानीय पड़ोस में रहती हैं। यह व्याकरण और स्थानीय अर्थ को सटीक रखता है।
लेन 2 (SA): "टेलीपोर्ट" लेन जहाँ कारें महत्वपूर्ण संदर्भ (context) प्राप्त करने के लिए हाईवे के दूरस्थ हिस्सों में रैंडमली कूद जाती हैं।
द गेट (The Gate): एक स्मार्ट ट्रैफिक कंट्रोलर (एक सीखा हुआ गेट) यह तय करता है कि प्रत्येक शब्द के लिए स्थानीय लेन बनाम टेलीपोर्ट लेन से कितनी जानकारी ली जानी चाहिए।

प्रयोगों ने क्या दिखाया

शोधकर्ताओं ने इसे दो तरीकों से परखा:

शून्य से एक नया AI बनाना: उन्होंने इस पद्धति का उपयोग करके एक नया लैंग्वेज मॉडल प्रशिक्षित किया। परिणाम? जिस मॉडल ने "शफल + लोकल" कॉम्बो का उपयोग किया, वह सबसे स्मार्ट था। इसने स्थानीय वाक्यों को अच्छी तरह समझा और फुल-अटेंशन मॉडल की तुलना में लंबी दूरी के संदर्भ को बेहतर ढंग से याद रखा।
मौजूदा AI (Qwen3) को अपग्रेड करना: उन्होंने एक शक्तिशाली, प्री-ट्रेंड AI (Qwen3) को लिया और बिना दोबारा ट्रेनिंग दिए, बस इसके अटेंशन मैकेनिज्म को इस नए "शफल" तरीके से बदल दिया।
- परिणाम: अपग्रेड किया गया AI लंबे टेक्स्ट को देखते समय मूल मॉडल की तुलना में बहुत बेहतर प्रदर्शन करता था। इसने एक फुल-अटेंशन मॉडल (जो एक साथ सब कुछ देखता है) की "बुद्धिमत्ता" को वापस पा लिया, लेकिन यह बहुत तेज़ी से चला और कम मेमोरी का उपयोग किया।

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

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

गति (Speed): यह तेज़ "लोकल विंडो" विधि की गति को बनाए रखता है।
बुद्धिमत्ता (Smarts): यह धीमे "सब कुछ देखने वाले" (look at everything) विधि की बुद्धिमत्ता प्राप्त करता है।
दक्षता (Efficiency): यह प्रकृति के सबसे कुशल नेटवर्क: फल मक्खी के मस्तिष्क की नकल करके इसे प्राप्त करता है।

संक्षेप में, यह शोध पत्र सिद्ध करता है कि सब कुछ जोड़ने के लिए आपको एक विशाल, महंगी नेटवर्क बनाने की आवश्यकता नहीं है। कभी-कभी, आपको बस चीजों को थोड़ा हिलाने (shake up) की आवश्यकता होती है।

तकनीकी सारांश: स्टोकेस्टिक अटेंशन (Stochastic Attention)

समस्या विवरण

ट्रांसफॉर्मर मॉडल्स को स्केल करने के लिए कुशल अटेंशन मैकेनिज्म महत्वपूर्ण हैं, फिर भी मौजूदा दृष्टिकोण कम्प्यूटेशनल लागत और रिसेप्टिव फील्ड कवरेज के बीच एक ट्रेड-ऑफ का सामना करते हैं। स्लाइडिंग-विंडो अटेंशन (SWA) प्रत्येक लेयर के लिए जटिलता को $O(nw) $तक कम कर देता है, लेकिन यह टोकन को आकार$ w$ के एक स्थानीय विंडो तक सीमित कर देता है। हालाँकि, यह नियतात्मक (deterministic) स्थानीयता प्रभावी रिसेप्टिव फील्ड को रैखिक विकास ( $\ell$ लेयर्स के बाद $\ell w$ ) तक सीमित कर देती है, जिससे अनुक्रम (sequence) का बड़ा हिस्सा अप्राप्य रह जाता है जब $w \ll n$ हो। जबकि ग्लोबल टोकन्स, हस्तनिर्मित स्पार्स पैटर्न, या ब्लॉक-स्तरीय रूटिंग जैसे उपाय मौजूद हैं, वे अक्सर आर्किटेक्चरल जटिलता बढ़ाते हैं या कवरेज गारंटी प्रदान करने में विफल रहते हैं। इसके विपरीत, फुल अटेंशन ग्लोबल कनेक्टिविटी तो प्रदान करता है लेकिन अत्यधिक $O(n^2)$ लागत वसूलता है।

यह शोध पत्र ड्रोसोफिला मेलानोगैस्टर (Drosophila melanogaster - फ्रूट फ्लाई) के कनेक्टोम से प्रेरणा लेता है, जो लगभग 130,000 न्यूरॉन्स का एक नेटवर्क है जिसमें कनेक्शन की संभावना केवल 0.02% है। इस स्पर्सिटी (sparsity) के बावजूद, नेटवर्क लगभग 4.4 हॉप्स की एक छोटी औसत पाथ लेंथ प्राप्त करता है। यह "स्मॉल-वर्ल्ड" टोपोलॉजी सघन स्थानीय क्लस्टरिंग और व्यापक रूप से वितरित लॉन्ग-रेंज कनेक्शनों के बीच के तालमेल से उत्पन्न होती है, जो स्टोकेस्टिक शॉर्टकट (stochastic shortcuts) के रूप में कार्य करते हैं। लेखक तर्क देते हैं कि कुशल वैश्विक सूचना प्रवाह, सघन ऑल-टू-ऑल कनेक्टिविटी की आवश्यकता के बिना, संरचित स्थानीय गणना और वितरित, स्टोकेस्टिक लॉन्ग-रेंज कनेक्शनों के माध्यम से उभर सकता है।

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

लेखक स्टोकेस्टिक अटेंशन (SA) का प्रस्ताव करते हैं, जो SWA के लिए एक पैरामीटर-मुक्त संवर्धन (enhancement) है जो कनेक्टोम की रूटिंग रणनीति की नकल करता है।

मुख्य तंत्र

SA एक निश्चित स्थानीय विंडो को तीन चरणों वाली प्रक्रिया के माध्यम से एक स्टोकेस्टिक ग्लोबल विंडो में बदल देता है, जिसे मानक SWA गणना से पहले और बाद में लागू किया जाता है:

परम्यूट (Permute): प्रत्येक लेयर के लिए स्वतंत्र रूप से एक रैंडम परम्यूटेशन $\sigma \sim \text{Uniform}(S_n)$ खींचा जाता है। टोकन सीक्वेंस रिप्रेजेंटेशन ( $Q, K, V$ ) को $\sigma$ के अनुसार पुनर्व्यवस्थित किया जाता है।
विंडोड अटेंशन (Windowed Attention): परम्यूटेड स्पेस में मानक SWA लागू किया जाता है। चूंकि परम्यूटेशन रैंडम है, इसलिए आकार $w$ की निश्चित विंडो अब मूल अनुक्रम के एक रैंडम सबसेट को कवर करती है।
अनडू परम्यूट (Undo Permute): आउटपुट पर $\sigma^{-1}$ लागू करके मूल क्रम को बहाल किया जाता है।

मूल टोकन स्पेस में, स्थिति $i$ एक अपेक्षित आकार $\tilde{N}^\sigma_w(i)$ के रैंडम पड़ोस पर ध्यान केंद्रित करती है, जो मूल दूरी $|i-j|$ की परवाह किए बिना पूरे अनुक्रम में समान रूप से वितरित होता है। यह किसी भी टोकन जोड़ी के लिए लगभग $w/n$ की पेयरवाइज कनेक्शन प्रोबेबिलिटी बनाता है, जो वितरित लॉन्ग-रेंज शॉर्टकट के समान है।

सैद्धांतिक गुण

रिसेप्टिव फील्ड ग्रोथ: जहाँ SWA का रिसेप्टिव फील्ड रैखिक रूप से बढ़ता है ( $O(n/w)$ लेयर्स पूर्ण कवरेज के लिए), वहीं SA के स्वतंत्र रूप से सैंपल किए गए परम्यूटेशन लेयर्स एक्सपोनेंशियल (घातांकीय) बढ़ते रिसेप्टिव फील्ड प्रदान करते हैं। पेपर सिद्ध करता है कि पूर्ण अनुक्रम कवरेज $O(\log_w n)$ लेयर्स में प्राप्त किया जा जाता है।
जटिलता: परम्यूटेशन और इनवर्स परम्यूटेशन ऑपरेशन $O(n)$ इंडेक्स पुनर्व्यवस्था हैं। अटेंशन गणना $O(nw)$ बनी रहती है। इस प्रकार, SA मानक SWA की तरह ही प्रति-लेयर $O(nw)$ बजट बनाए रखता है।
कॉज़ल कंस्ट्रेंट्स (Causal Constraints): ऑटोरेग्रेसिव सेटिंग्स में, कॉज़ल मास्क को परम्यूटेड स्पेस में लागू किया जाता है, जिससे यह सुनिश्चित होता है कि टोकन केवल कॉज़ली एक्सेसिबल पूर्ववर्तियों (predecessors) पर ध्यान दें, जबकि लॉन्ग-रेंज स्टोकेस्टिक कनेक्शन सक्षम रहते हैं।

गेटेड SA + SWA संयोजन

"स्मॉल-वर्ल्ड" रिजीम (उच्च स्थानीय क्लस्टरिंग + लघु पथ) को पुनरुत्पादित करने के लिए, लेखक SA और SWA के संयोजन वाला एक ड्यूल-पाथ आर्किटेक्चर प्रस्तावित करते हैं:
$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$
जहाँ $Y_{sa}$ और $Y_{swa}$ क्रमशः स्टोकेस्टिक और स्लाइडिंग-विंडो पाथ के आउटपुट हैं। गेटिंग वेट्स ( $g_{sa}, g_{swa}$ ) स्वतंत्र सिग्मॉइड प्रोजेक्शन के माध्यम से कंप्यूट किए जाते हैं, जिससे मॉडल को लोकल कोहेरेंस (SWA से) और ग्लोबल कवरेज (SA से) के बीच संतुलन बनाने की अनुमति मिलती है। यह संयोजन नगण्य पैरामीटर्स जोड़ता है और $O(nw)$ जटिलता बनाए रखता है।

मुख्य योगदान

स्टोकेस्टिक अटेंशन (SA): SWA के लिए एक ड्रॉप-इन, पैरामीटर-मुक्त संवर्धन जो $O(\log_w n)$ एक्सपोनेंशियल रिसेप्टिव फील्ड ग्रोथ प्राप्त करने के लिए रैंडम परम्यूटेशन लागू करता है, और यह भी उसी $O(nw)$ कम्प्यूटेशनल बजट के भीतर है।
गेटेड SA + SWAs: एक ड्यूल-पाथ आर्किटेक्चर जो जैविक कनेक्टोम में देखे गए स्मॉल-वर्ल्ड रिजीम को सैद्धांतिक और अनुभवजन्य रूप से पुनरुत्पादित करता है, जो स्थानीय क्लस्टरिंग और स्टोकेस्टिक शॉर्टकट के बीच संतुलन बनाता है।
अनुभवजन्य सत्यापन (Empirical Validation):
- प्री-ट्रेनिंग: स्क्रैच से प्रशिक्षित 360M-पैरामीटर मॉडल्स पर, गेटेड SA + SWA संयोजन सबसे अच्छा औसत ज़ीरो-शॉट सटीकता प्राप्त करता है, जो शुद्ध SWA से बेहतर और फुल अटेंशन की पर्प्लेक्सिटी (perplexity) के बराबर है, जबकि डाउनस्ट्रीम टास्क में उससे आगे निकल जाता है।
- ट्रेनिंग-फ्री इन्फरेंस: प्री-ट्रेंड Qwen3-8B और Qwen3-30B-A3B मॉडल्स पर लागू होने पर, SA लगातार SWA से बेहतर प्रदर्शन करता है और तुलनीय कंप्यूट बजट पर MoBA (Mixture of Block Attention) के बराबर या उससे अधिक प्रदर्शन करता है। यह विंडो साइज बढ़ने के साथ फुल-अटेंशन क्वालिटी को काफी तेजी से रिकवर करता है।

परिणाम

प्री-ट्रेनिंग: SA + SWA मॉडल ने 35.9 की औसत ज़ीरो-शॉट सटीकता प्राप्त की (SWA के 35.1 और फुल अटेंशन के 34.9 के मुकाबले)। इसने LAMBADA (जिसके लिए व्यापक डिस्कोर्स कॉन्टेक्स्ट की आवश्यकता होती है) पर प्रदर्शन में उल्लेखनीय सुधार किया, जबकि WikiText पर कम पर्प्लेक्सिटी बनाए रखा। शुद्ध SA उच्च पर्प्लेक्सिटी से ग्रस्त रहा, जो लैंग्वेज मॉडलिंग के लिए लोकल कोहेरेंस की आवश्यकता की पुष्टि करता है।
Qwen3 पर इन्फरेंस:
- 128 के प्रभावी विंडो साइज पर, Qwen3-8B पर SA ने 70.9% औसत सटीकता प्राप्त की (फुल-अटेंशन बेसलाइन के 1 पॉइंट के भीतर), जबकि SWA 62.2% पर पिछड़ गया।
- Qwen3-30B-A3B पर, SA ने 64 के विंडो साइज पर 73.2% सटीकता प्राप्त की, जो SWA (47.0%) और MoBA (66.3%) से काफी बेहतर है।
- SA ने बहुत छोटे विंडो साइज (जैसे $w=32$ ) पर भी नॉलेज-इंटेंसिव टास्क (MMLU, BoolQ) पर मजबूती प्रदर्शित की, जहाँ SWA का प्रदर्शन गिर गया था।
दक्षता (Efficiency): SA रैखिक स्केलिंग $O(nw)$ बनाए रखता है। एक A100 GPU पर, SA ने अनुक्रम लंबाई 2,048 पर फुल अटेंशन की तुलना में 1.5× की स्पीडअप दिखाई, जो 32,768 पर बढ़कर 28× हो गई।

महत्व और दावे

पेपर का दावा है कि कनेक्टोम-प्रेरित स्टोकेस्टिक रूटिंग कुशल अटेंशन की एक्सप्रेसिविटी को सुधारने के लिए एक व्यावहारिक प्रिमिटिव (primitive) है। लेखक तर्क देते हैं कि ग्लोबल सूचना प्रवाह के लिए सघन ऑल-टू-ऑल कनेक्टिविटी की आवश्यकता नहीं है, बल्कि यह डेप्थ के साथ संचित स्ट्रक्चर्ड लोकल कंप्यूटेशन और स्पार्स, स्टोकेस्टिक लॉन्ग-रेंज शॉर्टकट के तालमेल से उभर सकता है।

इस कार्य का महत्व निम्नलिखित है:

सरलता: SA को अटेंशन मैकेनिज्म में किसी भी सीखने योग्य (learnable) पैरामीटर की आवश्यकता नहीं है और इसमें केवल $O(n)$ इंडेक्स ओवरहेड है, जो इसे मौजूदा SWA कर्नेल्स (जैसे Mistral, Gemma 2, gpt-oss में) के लिए एक "ड्रॉप-इन" अपग्रेड बनाता है।
सैद्धांतिक गारंटी: यह $O(\log_w n)$ कवरेज डेप्थ प्रदान करता है, जो SWA के रैखिक कवरेज की तुलना में एक महत्वपूर्ण सुधार है।
पूरकता: इसे मौजूदा लीनियर और स्पार्स दृष्टिकोणों के पूरक के रूप में प्रस्तुत किया गया है, जो एक नया डिज़ाइन सिद्धांत प्रदान करता है जो हैंड-क्राफ्टेड स्पार्स पैटर्न या ग्लोबल टोकन्स की जटिलता के बिना लोकल एफिशिएंसी और ग्लोबल एक्सप्रेसिविटी के बीच के अंतर को पाटता है।