💬 NLP

DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention

DashAttention, $\alpha$ -entmax का उपयोग करके परिवर्तनीय संख्या में KV ब्लॉक्स को गतिशील रूप से चुनने के लिए एक पूर्णतः विभेदनीय (differentiable) और अनुकूलन योग्य विरल पदानुक्रमित अटेंशन (sparse hierarchical attention) तंत्र पेश करता है, जो NSA और InfLLMv2 जैसी मौजूदा विधियों की तुलना में बेहतर लॉन्ग-कॉन्टेक्स्ट मॉडलिंग सटीकता और इन्फरेंस गति प्राप्त करता है।

मूल लेखक: Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M. Ponti, André F. T. Martins, Marcos V. Treviso

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

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

CC BY 4.0

मूल लेखक: Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M. Ponti, André F. T. Martins, Marcos V. Treviso

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

कल्पना कीजिए कि आप एक अकेले प्रश्न का उत्तर खोजने के लिए एक विशाल, 1,00,000 पृष्ठों के विश्वकोश (encyclopedia) को पढ़ने की कोशिश कर रहे हैं।

समस्या: "सब कुछ" बनाम "Top-K" की दुविधा
वर्तमान AI मॉडल (लार्ज लैंग्वेज मॉडल्स) आमतौर पर इसे दो तरीकों से संभालते हैं, जिनमें दोनों की अपनी कमियां हैं:

"सब कुछ पढ़ने" वाला दृष्टिकोण (Full Attention): मॉडल उत्तर खोजने के लिए विश्वकोश के हर एक शब्द को पढ़ने की कोशिश करता है। यह सटीक तो है लेकिन अविश्वसनीय रूप से धीमा और महंगा है, जैसे केवल एक रेसिपी खोजने के लिए पूरी किताब पढ़ डालना।
"टॉप 5 चुनने" वाला दृष्टिकोण (Top-K Sparse Attention): मॉडल तेजी से विषय-सूची (table of contents) को स्कैन करता है, उन शीर्ष 5 अध्यायों को चुनता है जो उसे लगता है कि प्रासंगिक हैं, और बाकी को अनदेखा कर देता है। यह तेज़ है, लेकिन कठोर है। क्या होगा यदि उत्तर वास्तव में अध्याय 6 में है? या क्या होगा यदि उत्तर के लिए 20 अलग-अलग बिखरे हुए पृष्ठों को पढ़ने की आवश्यकता है? साथ भी, एक बार जब मॉडल उन 5 अध्यायों को चुन लेता है, तो वह उन अध्यायों से "सीख" नहीं सकता जिन्हें उसने अनदेखा कर दिया था, जिससे प्रशिक्षण (training) की प्रक्रिया जटिल हो जाती है।

समाधान: DashAttention
इस शोध पत्र के लेखक DashAttention नामक एक नई विधि प्रस्तावित करते हैं। DashAttention को एक स्मार्ट, अनुकूलन योग्य (adaptive) लाइब्रेरियन के रूप में समझें जो केवल निश्चित संख्या में किताबें नहीं चुनता, बल्कि यह तय करता है कि प्रश्न की जटिलता के आधार पर उसे कितनी किताबों को पढ़ने की आवश्यकता है।

DashAttention इस प्रकार कार्य करता है, जिसे तीन चरणों में एक सरल उपमा (analogy) के माध्यम से समझाया गया है:

चरण 0: "अध्याय सारांश" (Local Chunk Summarization)

तुरंत हर एक शब्द को देखने के बजाय, मॉडल पहले विशाल पाठ को छोटे "चंक्स" (जैसे अध्यायों) में तोड़ देता है।

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

चरण 1: "अनुकूलन योग्य गेटकीपर" (Entmax Routing)

अब, मॉडल के पास अध्यायों के सारांशों की एक सूची है। उसे यह तय करना है कि किन अध्यायों को विस्तार से पढ़ना है।

पुराना तरीका (Top-K): मॉडल का एक सख्त नियम होता है: "हमेशा ठीक 5 अध्याय चुनें।" यदि प्रश्न सरल है, तो वह 5 अध्याय पढ़ने में समय बर्बाद करता है। यदि प्रश्न कठिन है, तो वह महत्वपूर्ण जानकारी खो देता है क्योंकि वह 5 पर सीमित है।
DashAttention का तरीका: मॉडल $\alpha$ -entmax नामक एक विशेष गणितीय उपकरण का उपयोग करता है। एक गेटकीपर की कल्पना करें जो प्रश्न और सारांशों को देखता है।
- यदि प्रश्न सरल है ("फ्रांस की राजधानी क्या है?"), तो गेटकीपर कहता है, "केवल 1 अध्याय की आवश्यकता है," और बाकी को बाहर कर देता है।
- यदि प्रश्न जटिल है ("तीन महाद्वीपों में व्यापार मार्गों के इतिहास का पता लगाएं"), तो गेटकीपर कहता है, "ठीक है, हमें 15 अध्यायों की आवश्यकता है," और द्वार को अधिक खोल देता है।
- जादू: यह गेटकीपर "डिफरेंशिएबल" (differentiable) है। इसका अर्थ है कि मॉडल एक बेहतर गेटकीपर कैसे बनना है, यह सीख सकता है। यदि यह प्रशिक्षण के दौरान गलत अध्याय चुनता है, तो इसे अपनी गेटकीपिंग रणनीति को समायोजित करने के लिए एक संकेत मिलता है। यह कोई सख्त "हाँ/नहीं" वाला स्विच नहीं है; यह एक सुचारू, सीखने योग्य डायल है।

चरण 2: "गहन अध्ययन" (Prior-Induced Sparse Softmax)

अंत में, मॉडल गेटकीपर द्वारा चयनित विशिष्ट अध्यायों को पढ़ता है।

यह गेटकीपर से मिले "मतों" (votes) को लेता है और चुनी गई सामग्री के विस्तृत अध्ययन के लिए उनका उपयोग करता है।
यह सुनिश्चित करता है कि भले ही इसने अधिकांश पुस्तक को छोड़ दिया हो, फिर भी यह कहानी के प्रवाह को न खोए। यह अंतराल को इस तरह भरता है कि अंतिम उत्तर उतना ही सटीक हो जितना कि पूरी किताब पढ़ने से होता, लेकिन इसने यह बहुत तेज़ी से किया।

यह बेहतर क्यों है? (परिणाम)

शोध पत्र का दावा है कि DashAttention तीन प्रमुख क्षेत्रों में जीतता है:

स्मार्ट चयन: कठोर "टॉप-5" नियम के विपरीत, DashAttention अनुकूलित होता है। यह कठिन प्रश्नों पर अधिक "मस्तिष्क शक्ति" खर्च करता है और आसान प्रश्नों पर कम। यह इसे ढेर में से विशिष्ट सुई खोजने (retrieval tasks) में बहुत बेहतर बनाता है।
कोई "बिखराव" नहीं (No Dispersion): लंबे टेक्स्ट में, मानक AI मॉडल अक्सर "विचलित" हो जाते हैं और अपना ध्यान बहुत अधिक फैला देते हैं, जैसे कि टॉर्च की रोशनी का बहुत चौड़ा होना जिससे कुछ भी स्पष्ट दिखाई न दे। DashAttention ध्यान को केंद्रित रखता है, यह सुनिश्चित करता है कि भारी मात्रा में टेक्स्ट होने पर भी मॉडल सटीक बना रहे।
गति: क्योंकि यह अप्रासंगिक हिस्सों को पढ़ने से बच जाता है, इसलिए यह अविश्वसनीय रूप से तेज़ है।
- लेखकों ने कंप्यूटर चिप्स (GPUs) के लिए एक विशेष संस्करण बनाया है जो बहुत लंबे टेक्स्ट के साथ काम करते समय वर्तमान उद्योग मानक (FlashAttention-3) की तुलना में 3.36 गुना तेज़ चलता है।
- यह वही सटीकता प्राप्त करता है जो पूरी किताब पढ़ने से मिलती है, लेकिन केवल 25% कंप्यूटिंग पावर (75% स्पर्सिटी) का उपयोग करता है।

सारांश

DashAttention एक कठोर, नियम-बद्ध लाइब्रेरियन से एक अत्यधिक बुद्धिमान, अनुकूलन योग्य सहायक में अपग्रेड करने जैसा है जो हमेशा 5 किताबें चुनता है, बल्कि एक ऐसा सहायक है जो विषय-सूची पढ़ता है, यह तय करता है कि विशिष्ट प्रश्न के लिए कितने अध्यायों की आवश्यकता है, और फिर केवल उन्हीं में गहराई से उतरता है। यह पिछले तरीकों की तुलना में तेज़, स्मार्ट और बेहतर सीखने वाला है, जिससे यह संभव हो जाता है कि AI भारी मात्रा में जानकारी को बिना अभिभूत हुए या धीमे हुए संभाल सके।

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

समस्या विवरण

वर्तमान पदानुक्रमित (hierarchical) अटेंशन विधियाँ, जैसे कि NSA और InfLLMv2, मोटे तौर पर (coarse) अटेंशन स्कोर के आधार पर शीर्ष- $k$ प्रासंगिक की-वैल्यू (KV) ब्लॉक्स का चयन करके लंबे संदर्भ मॉडलिंग (long-context modeling) की कम्प्यूटेशनल बाधाओं को संबोधित करती हैं। हालाँकि, ये दृष्टिकोण दो प्राथमिक सीमाओं से ग्रस्त हैं:

निश्चित बजट और लचीलेपन का अभाव: शीर्ष- $k$ ऑपरेशन किसी भी क्वेरी के लिए एक निश्चित संख्या में प्रासंगिक टोकन मान लेता है, जो विभिन्न टोकन, अटेंशन हेड्स और लेयर्स के बीच अलग-अलग सूचनात्मक मांगों के अनुकूल होने में विफल रहता है।
अ-विभेदकता (Non-Differentiability): हार्ड टॉप- $k$ चयन, मोटे रूटिंग चरण और सूक्ष्म (fine-grained) टोकन-स्तरीय अटेंशन के बीच ग्रेडिएंट प्रवाह को बाधित कर देता है। यह मॉडल को यह सीखने से रोकता है कि ब्लॉक सदस्यता बदलने के लिए रूटिंग स्कोर को कैसे समायोजित किया जाए, जिससे स्पार्स चयन और अंतिम लॉस (loss) के बीच एक विच्छेद पैदा होता है।
प्रसार (Dispersion): हालिया सैद्धांतिक कार्य सुझाव देते हैं कि लंबे संदर्भों में मानक सॉफ्टमैक्स अटेंशन "प्रसारशील" (dispersive) हो सकता है, जहाँ अटेंशन वितरण की शैनन एंट्रॉपी (Shannon entropy) अनुक्रम लंबाई के साथ लॉगरिदमिक रूप से बढ़ती है, जिससे लंबी-दूरी की निर्भरता मॉडलिंग कठिन हो जाती है। जबकि टॉप- $k$ इसे कम करता है, मौजूदा पदानुक्रमित विधियाँ जो चयन से पहले हेड्स को सॉफ्टमैक्स के माध्यम से एकत्रित करती हैं, प्रसार को पुन: उत्पन्न कर सकती हैं।

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

DashAttention (डिफरेंशिएबल और एडेप्टिव स्पार्स हिरार्किकल अटेंशन) एक बहु-चरणीय तंत्र है जिसे पूरी तरह से डिफरेंशिएबल और एडेप्टिवली स्पार्स होने के लिए डिज़ाइन किया गया है। यह हार्ड टॉप- $k$ रूटिंग को $\alpha$ -entmax ट्रांसफॉर्मेशन से बदल देता है, जिससे मॉडल इनपुट के आधार पर गतिशील रूप से प्रासंगिक ब्लॉक्स की संख्या सीख सकता है।

इसकी वास्तुकला तीन चरणों से बनी है:

चरण 0: स्थानीय चंक सारांश (Local Chunk Summarization)

संदर्भ को आकार $B$ के निरंतर चंक्स (chunks) में विभाजित किया जाता है। साधारण मीन पूलिंग या सीखे गए MLP (जो प्रीट्रेन्ड मॉडल्स से अनुकूलित करना कठिन हो सकता है) के बजाय, DashAttention एक सीखा गया स्थानीय अटेंशन तंत्र उपयोग करता है। एक सीखा गया क्वेरी वेक्टर $\bar{q}$ लोकल स्केल्ड डॉट-प्रोडक्ट अटेंशन (SDPA) के माध्यम से प्रत्येक चंक के भीतर कीज़ (keys) पर ध्यान केंद्रित करता है।

आरंभीकरण (Initialization): $\bar{q}$ को शून्य पर आरंभीकृत किया जाता है, जिससे प्रशिक्षण की शुरुआत में यह तंत्र यूनिफॉर्म मीन पूलिंग की तरह व्यवहार करता है।
अनुकूलन (Adaptation): जैसे-जैसे प्रशिक्षण आगे बढ़ता है, $\bar{q}$ चंक के भीतर टोकन को वेट देने के लिए सीखता है, जिससे अधिक अभिव्यंजक सारांश (expressive summaries) बनते हैं।
दक्षता: एक बार जब एक चंक पूरी तरह से जेनरेट हो जाता है, तो उसका सारांश स्थिर और कैश (cache) कर दिया जाता है, जिससे इन्फरेंस के दौरान पुन: गणना सुनिश्चित नहीं होती है।

चरण 1: Entmax चंक रूटिंग

प्रत्येक क्वेरी के लिए, मॉडल चंक सारांशों के विरुद्ध लॉजिट्स (logits) की गणना करता है। टॉप- $k$ ट्रंकेशन लागू करने के बजाय, इन लॉजिट्स को $\alpha$ -entmax का उपयोग करके रूपांतरित किया जाता है:
$\hat{w} = \alpha\text{-entmax}(\gamma \bar{z})$

एडेप्टिव स्पर्सिटी (Adaptive Sparsity): $\alpha$ -entmax एक स्पार्स प्रोबेबिलिटी वेक्टर उत्पन्न करता है जहाँ गैर-शून्य तत्वों (support) की संख्या इनपुट स्कोर द्वारा निर्धारित होती है, न कि किसी निश्चित हाइपरपैरामीटर द्वारा।
विभेदकता (Differentiability): टॉप- $k$ के विपरीत, $\alpha$ -entmax पूरी तरह से डिफरेंशिएबल है। यह ग्रेडिएंट को चंक सारांश चरण (चरण 0) तक वापस प्रवाहित होने की अनुमति देता है।
GQA हैंडलिंग: ग्रुपड क्वेरी अटेंशन (GQA) सेटिंग्स में, एक साझा KV हेड साझा करने वाले क्वेरी हेड्स के लिए रूटिंग प्रोबेबिलिटी को एक साझा स्पार्स सपोर्ट बनाने के लिए औसत निकाला जाता है।

चरण 2: प्रायर-इंड्यूस्ड स्पार्स सॉफ्टमैक्स (Prior-Induced Sparse Softmax)

अंतिम चरण टोकन स्तर पर अटेंशन को परिष्कृत करता है लेकिन केवल उन्हीं चंक्स के भीतर जो चरण 1 में चुने गए हैं। डिफरेंशिएबिलिटी बनाए रखने और मोटे रूटिंग से सूक्ष्म अटेंशन के संक्रमण को संभालने के लिए, यह विधि एक प्रायर-इंड्यूस्ड सॉफ्टमैक्स पेश करती है:

वेरिएशनल फॉर्मूलेशन: मानक सॉफ्टमैक्स को $z^\top p - KL(p \| u)$ को मैक्सिमाइज करने के रूप में देखा जाता है। DashAttention संदर्भ वितरण $u$ को चरण 1 के रूटिंग स्कोर $w$ से प्राप्त वितरण $g$ में बदल देता है।
डायनेमिक मास्किंग: $g$ का निर्माण इस प्रकार किया जाता है कि जब भी रूटिंग स्कोर $w_{i,j} = 0$ हो, तो $g_{i,j} = 0$ हो जाए, जिसके परिणामस्वरूप सॉफ्टमैक्स अटेंशन स्वाभाविक रूप से स्पार्स हो जाता है। शून्य-प्रोबेबिलिटी वाले चंक्स प्रभावी रूप से मास्क कर दिए जाते हैं।
डायगोनल हैंडलिंग: डायगोनल के पास स्थित चंक्स (जिनके पास सारांश नहीं हो सकता) के लिए, एक अलग ब्रांच बनाई जाती है। एक हाइपरपैरामीटर $\sigma$ प्रायर की शक्ति को नियंत्रित करता है, जिससे मॉडल को रूटिंग स्कोर और एक यूनिफॉर्म बेसलाइन के बीच संतुलन बनाने की अनुमति मिलती है।
कार्यान्वयन: इस चरण को एक संशोधित फ्लैश-अटेंशन कर्नेल का उपयोग करके लागू किया गया है जो बिट-पैक्ड मास्क द्वारा परिभाषित सक्रिय चंक्स पर इटरेट करता है, और सॉफ्टमैक्स से पहले लॉजिट्स में रूटिंग बायस जोड़ता है।

मुख्य योगदान

डिफरेंशिएबल हिरार्किकल अटेंशन: DashAttention एक एंड-टू-एंड डिफरेंशिएबल स्पार्स अटेंशन तंत्र पेश करता है। $\alpha$ -entmax का उपयोग करके, यह टॉप- $k$ विधियों में निहित ग्रेडिएंट ब्लॉकेज को समाप्त करता है, जिससे मॉडल इष्टतम स्पर्सिटी पैटर्न सीख सकता है।
एडेप्टिव स्पर्सिटी एलोकेशन: निश्चित-बजट विधियों के विपरीत, DashAttention टोकन, हेड्स और लेयर्स के बीच गतिशील रूप से स्पर्सिटी आवंटित करता है। सूचनात्मक क्वेरी अधिक ब्लॉक्स पर ध्यान दे सकती हैं, जबकि तीव्र (sharp) क्वेरी कम ब्लॉक्स पर, जिससे प्रति इंस्टेंस कंप्यूट बजट अनुकूलित होता है।
गैर-प्रसार गुण (Non-Dispersive Property): सैद्धांतिक विश्लेषण दर्शाता है कि हेड्स के एकत्रीकरण के लिए $\alpha$ -entmax का उपयोग करना प्रसार की घटना को रोकता है जो सॉफ्टमैक्स-आधारित हिरार्किकल अटेंशन में सामान्य है, जो सैद्धांतिक रूप से लंबे संदर्भ मॉडलिंग क्षमताओं में सुधार करता है।
कुशल GPU कार्यान्वयन: लेखकों ने तीन फ्यूज्ड कर्नेल्स (प्रत्येक चरण के लिए एक) के साथ एक ट्रिटन-आधारित (Triton-based) कार्यान्वयन प्रदान किया है। यह डिज़ाइन मौजूदा कुशल अटेंशन लाइब्रेरीज़ (जैसे फ्लैश-अटेंशन) के साथ संगत है और प्रीफिल एवं इंक्रीमेंटल डिकोडिंग दोनों का समर्थन करता है।

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

लेखकों ने 16K कॉन्टेक्स्ट लेंथ के साथ MiniCPM-4 मॉडल्स (1B, 3B, और 8B पैरामीटर्स) पर Full Attention, NSA और InfLLMv2 के मुकाबले DashAttention का मूल्यांकन किया।

टास्क प्रदर्शन: RULER और HELMET बेंचमार्क पर, DashAttention ने NSA और InfLLMv2 की तुलना में लगातार बेहतर प्रदर्शन किया, विशेष रूप से चुनौतीपूर्ण लंबे-संदर्भ रिट्रीवल कार्यों (जैसे MK1–MK3) में। 75% स्पर्सिटी पर, इसने Full Attention के बराबर सटीकता प्राप्त की।
सामान्यीकरण (Generalization): छोटे-संदर्भ वाले सामान्य कार्यों (MMLU, GSM8K, आदि) पर, DashAttention ने Full Attention मॉडल्स के समान प्रदर्शन किया, जो सामान्य क्षमताओं में कोई महत्वपूर्ण गिरावट नहीं दिखाता है।
दक्षता:
- प्रीफिल (Prefill): DashAttention ने FlashAttention-3 की तुलना में 1.34× से 3.09× तक की गति वृद्धि (speedup) प्राप्त की।
- डिकोडिंग (Decoding): मेमोरी-बाउंड रिजीम में, इसने FlashAttention-3 पर 3.36× और InfLLMv2 पर 1.35× तक की गति वृद्धि प्राप्त की।
पारेटो फ्रंटियर (Pareto Frontier): DashAttention ने सटीकता-स्पर्सिटी ट्रेड-ऑफ कर्व पर प्रभुत्व जमाया, उच्च स्पर्सिटी स्तरों (जैसे ~90% स्पर्सिटी) पर भी उच्च सटीकता बनाए रखी।
गतिशील व्यवहार: एब्लेशन अध्ययनों ने पुष्टि की कि DashAttention स्वाभाविक रूप से शुरुआती लेयर्स की तुलना में मध्य लेयर्स को अधिक स्पर्सिटी आवंटित करता है, जो बिना किसी स्पष्ट ट्यूनिंग के इष्टतम बजट आवंटन रणनीतियों की नकल करता है।

महत्व और दावे

पेपर का दावा है कि DashAttention स्पार्स अटेंशन (दक्षता) और डेंस अटेंशन (लचीलापन और डिफरेंशिएबिलिटी) के लाभों को एकीकृत करके लंबे संदर्भों को मॉडल करने के लिए एक लागत-प्रभावी रणनीति प्रदान करता है।

सैद्धांतिक लाभ: यह $\alpha$ -entmax एकत्रीकरण के माध्यम से यह सुनिश्चित करके "प्रसार" समस्या को संबोधित करता है कि अटेंशन तंत्र गैर-प्रसारशील बना रहे।
व्यावहारिक लाभ: यह एक पूर्णतः डिफरेंशिएबल, ट्रेन करने योग्य विधि प्रदान करके सैद्धांतिक स्पार्स अटेंशन और व्यावहारिक परिनियोजन (deployment) के बीच के अंतर को पाटता है, जिसे बिना किसी आर्किटेक्चरल ओवरहाल के प्रीट्रेन्ड सॉफ्टमैक्स मॉडल्स से अनुकूलित किया जा सकता है।
प्रदर्शन: यह विधि एक बेहतर पारेटो फ्रंटियर प्राप्त करती है, जो मौजूदा हिरार्किकल विधियों की तुलना में उच्च स्पर्सिटी स्तरों पर बेहतर सटीकता प्रदान करती है, साथ ही महत्वपूर्ण इन्फरेंस स्पीडअप भी देती है।

लेखक एक सीमा नोट करते हैं: कर्नेल्स को अभी तक vLLM या SGLang जैसे प्रमुख सर्विंग फ्रेमवर्क्स में एकीकृत नहीं किया गया है, जिसे भविष्य के कार्य के रूप में पहचाना गया है। वे हाइब्रिड मॉडल आर्किटेक्चर के लिए संभावित विस्तार का भी सुझाव देते हैं।