🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

AdaSplash-2 एक नवीन हिस्टोग्राम-आधारित इनिशियलाइज़ेशन और स्पैरसिटी-अवेयर (sparsity-aware) GPU कार्यान्वयन पेश करता है जो $\alpha$ -entmax स्पार्स अटेंशन को नाटकीय रूप से तेज़ करता है, जिससे मध्यम-से-उच्च स्पैरसिटी स्तरों पर FlashAttention-2 के तुलनीय प्रशिक्षण गति सक्षम होती है जबकि लॉन्ग-कॉन्टेक्स्ट सेटिंग्स में बेहतर प्रदर्शन प्राप्त होता है।

मूल लेखक: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

प्रकाशित 2026-04-17

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

CC BY 4.0

मूल लेखक: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

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

कल्पना कीजिए कि आप एक विशाल पुस्तकालय (एक ट्रांसफॉर्मर AI मॉडल) चला रहे हैं जहाँ हर किताब (वाक्य में एक शब्द) को संदर्भ समझने के लिए दूसरी हर किताब को पढ़ना पड़ता है।

एक पारंपरिक पुस्तकालय में, यह एक बुरा सपना है। यदि आपके पास 1,000 किताबें हैं, तो आपको 1,000,000 जोड़ों के कनेक्शनों को पढ़ना होगा। यदि आपके पास 100,000 किताबें हैं, तो कनेक्शनों की संख्या बढ़कर 10 अरब हो जाएगी। यह वह क्वाड्रेटिक बॉटलनेक (quadratic bottleneck) है जो आधुनिक AI को लंबे दस्तावेज़ों को पढ़ने में धीमा कर देता है।

इसे ठीक करने के लिए, शोधकर्ताओं ने स्पार्स अटेंशन (Sparse Attention) का आविष्कार किया। हर किताब को पढ़ने के बजाय, AI यह सीखता है कि अप्रासंगिक (irrelevant) किताबों को कैसे अनदेखा किया जाए और केवल महत्वपूर्ण किताबों पर ही ध्यान केंद्रित किया जाए। यह एक ऐसे लाइब्रेरियन की तरह है जो आपकी क्वेरी के लिए केवल 5 सबसे प्रासंगिक किताबें निकालता है, और बाकी को अनदेखा कर देता है।

हालाँकि, इसमें एक पेंच था। वे जिस "स्मार्ट लाइब्रेरियन" एल्गोरिदम का उपयोग कर रहे थे (जिसे $\alpha$ -entmax कहा जाता है), उसे यह तय करने में बहुत समय लग रहा था कि किन किताबों को चुनना है। इसे बहुत कठिन गणित करना पड़ता था, जैसे एक लाइब्रेरियन हर एक किताब के लिए सटीक संभावना (probability) की गणना करने की कोशिश कर रहा हो ताकि यह तय किया जा सके कि किसे अनदेखा करना है। यह गणना इतनी धीमी थी कि इसने किताबों को अनदेखा करने से होने वाले गति के लाभ को ही खत्म कर दिया।

मिलिए ADASPLASH-2 से: "सुपर-स्पीडी लाइब्रेरियन।"

यह पेपर इस स्मार्ट लाइब्रेरियन एल्गोरिदम को चलाने का एक नया, तेज़ तरीका पेश करता है। यह कैसे काम करता है, इसके लिए कुछ रोज़मर्रा के उदाहरणों का उपयोग किया गया है:

1. समस्या: "अनुमान लगाने का खेल" (The Guessing Game)

यह तय करने के लिए कि किन किताबों को अनदेखा किया जाए, एल्गोरिदम को एक विशिष्ट "कटऑफ स्कोर" (जिसे $\tau$ कहा जाता है) खोजना होता है।

पुराना तरीका: कल्पना करें कि आप पानी उबलने का सटीक तापमान खोजने की कोशिश कर रहे हैं। आपको बार-बार अनुमान लगाना पड़ता है, जांचना पड़ता है और अपने अनुमान को समायोजित करना पड़ता है। इसमें बहुत समय लगता है।
समस्या: पिछले संस्करण (ADASPLASH) में, कंप्यूटर को इस कटऑफ को खोजने के लिए कई बार सभी स्कोर्स को स्कैन करना पड़ता था। यह सटीक तो था लेकिन धीमा था।

2. समाधान: "हिस्टोग्राम शॉर्टकट" (The Histogram Shortcut)

ADASPLASH-2 के लेखकों ने एक शानदार तरकीब निकाली है: हिस्टोग्राम (Histogram)।

हर एक स्कोर को एक-एक करके देखने के बजाय, कल्पना करें कि आपके पास 100 स्लॉट्स (bins) वाला एक बड़ा बाल्टी जैसा डिब्बा है।

जैसे ही कंप्यूटर स्कोर्स को स्कैन करता है, वह सटीक नंबरों को स्टोर नहीं करता है। इसके बजाय, वह बस उस स्लॉट में एक कंचा (marble) डाल देता है जो उस स्कोर की रेंज से मेल खाता है।
जादू: जब तक वह स्कैनिंग पूरी करता है, उसके पास एक त्वरित "मैप" (हिस्टोग्राम) होता है कि सभी स्कोर कहाँ क्लस्टर (एकत्रित) हैं।
परिणाम: यह मैप कंप्यूटर को तुरंत कटऑफ स्कोर का एक बहुत अच्छा अनुमान दे देता है। यह कंचों के वितरण को देखकर यह कहने जैसा है कि, "ठीक है, उबाल बिंदु निश्चित रूप से 90 और 95 डिग्री के बीच है," बिना हर एक बूंद को मापने की आवश्यकता के।

चूंकि यह "मैप" कंप्यूटर की सुपर-फास्ट इंटरनल मेमोरी (SRAM) में तुरंत बनाया जाता है, इसलिए एल्गोरिदम को सटीक उत्तर खोजने के लिए दर्जनों बार के बजाय केवल 1 या 2 त्वरित जांचों की आवश्यकता होती है।

3. "खाली बक्सों को छोड़ें" वाला तरीका (The "Skip the Empty Boxes" Trick)

एक बार जब कंप्यूटर को कटऑफ पता चल जाता है, तो उसे पता चल जाता है कि कौन सी किताबें (या डेटा के ब्लॉक्स) अप्रासंगिक हैं।

पुराना तरीका: कंप्यूटर अभी भी खाली शेल्फ के पास से गुजरता, उन्हें चेक करता, यह महसूस करता कि वे खाली हैं, और फिर आगे बढ़ जाता। इससे ईंधन (ऊर्जा और समय) बर्बाद होता था।
नया तरीका: ADASPLASH-2 एक बिट-पैक्ड मास्क (bit-packed mask) का उपयोग करता है। इसे एक अत्यधिक कुशल चेकलिस्ट के रूप में सोचें जहाँ एक एकल "1" या "0" कंप्यूटर को ठीक से बताता है कि किन शेल्फों पर जाना है। यदि कोई शेल्फ खाली है, तो कंप्यूटर वास्तव में उस शेल्फ तक जाने की ड्राइव को स्किप (छोड़) कर देता है। वह उस डेटा को मेमोरी में लोड तक नहीं करता है।

यह क्यों मायने रखता है?

गति (Speed): ऐसी स्थितियों में जहाँ AI बहुत लंबे दस्तावेज़ (जैसे कि पूरा उपन्यास या कानूनी अनुबंध) पढ़ रहा है, "स्मार्ट लाइब्रेरियन" (ADASPLASH-2) अब मानक "डम्ब लाइब्रेरियन" (FlashAttention-2) से तेज़ है जो सब कुछ पढ़ता है।
स्मार्ट लॉन्ग-कॉन्टेक्स्ट (Smarter Long-Context): क्योंकि AI गणित में उलझे बिना लंबे टेक्स्ट को अधिक कुशलता से संभाल सकता है, इसलिए यह लंबे कार्यों में वास्तव में अधिक स्मार्ट हो जाता है। यह अप्रासंगिक विवरणों से भ्रमित होना बंद कर देता है और कहानी पर ध्यान केंद्रित करता है।
कोई समझौता नहीं (No Trade-off): आमतौर पर, चीजों को तेज़ बनाने से उनकी सटीकता कम हो जाती है। लेकिन यहाँ, नया तरीका इतना कुशल है कि AI छोटे कार्यों में उतना ही अच्छा प्रदर्शन करता है (या बेहतर करता है) और लंबे कार्यों में काफी बेहतर प्रदर्शन करता है।

निचोड़ (The Bottom Line)

ADASPLASH-2 एक लाइब्रेरियन को अपग्रेड करने जैसा है—जो उत्तर खोजने के लिए हर किताब के हर पन्ने को पढ़ने वाले व्यक्ति से, एक सुपर-इंटेलिजेंट असिस्टेंट में बदल जाता है जो:

मोटे तौर पर अनुमान लगाने के लिए स्पाइन (किताब के बाहरी भाग) को जल्दी से स्कैन करता है (The Histogram)।
तुरंत जानता है कि किन किताबों को शेल्फ से निकालना है।
खाली शेल्फों को पूरी तरह से अनदेखा कर देता है।

यह AI मॉडल्स को लंबी किताबें पढ़ने, जटिल कहानियों को समझने और यह सब पलक झपकते ही करने की अनुमति देता है।

यहाँ शोध पत्र "ADASPLASH-2: Faster Differentiable Sparse Attention" का विस्तृत तकनीकी सारांश दिया गया है।

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

ट्रांसफॉर्मर मॉडल, विशेष रूप से लंबे संदर्भ (long-context) वाले परिदृश्यों में, सेल्फ-अटेंशन मैकेनिज्म की द्विघातीय (quadratic) समय और मेमोरी जटिलता ( $O(n^2)$ ) के कारण एक कम्प्यूटेशनल बाधा का सामना करते हैं। हालाँकि FlashAttention (और इसके उत्तराधिकारी) ने ऑपरेशन्स को फ्यूज करके और ऑन-चिप मेमोरी का उपयोग करके मानक softmax अटेंशन को अनुकूलित किया है, लेकिन softmax स्वाभाविक रूप से सघन (dense) रहता है (प्रत्येक टोकन को गैर-शून्य प्रायिकता प्राप्त होती है)। यह घनत्व लंबे संदर्भों में अप्रासंगिक टोकनों पर अटेंशन मास को फैला देता है, जिससे रिप्रजेंटेशनल कोलैप्स (representational collapse) और खराब सामान्यीकरण (generalization) होता है।

$\alpha$ -entmax अटेंशन एक आशाजनक विकल्प है जो इनपुट-डिपेंडेंट स्पार्स (sparse) आउटपुट (सटीक शून्य) उत्पन्न करता है और डिफरेंशिएबल भी रहता है। हालाँकि, $\tau$ के लिए नॉर्मलाइजेशन कांस्टेंट (normalization constant) की गणना करने के लिए $\alpha$ -entmax को पुनरावृत्ति (iterative) रूट-फाइंडिंग विधियों की आवश्यकता होती है (softmax के क्लोज्ड-फॉर्म लॉग-सम-एक्सप के विपरीत)। इसे GPU पर लागू करने के पिछले प्रयास (जैसे ADASPLH) $\tau$ पर अभिसरण (converge) करने के लिए आवश्यक कई पास (passes) के कारण उच्च कम्प्यूटेशनल ओवरहेड से ग्रस्त थे, जिससे वे FlashAttention-2 जैसे अनुकूलित सॉफ्टमैक्स कर्नेल्स की गति से मेल खाने में असमर्थ रहे।

2. कार्यप्रणाली: ADASPLASH-2 (Methodology: ADASPLASH-2)

लेखक ADASPLASH-2 का प्रस्ताव करते हैं, जो एक हार्डवेयर-अवेयर एल्गोरिदम है जो ऑन-चिप SRAM और हिस्टोग्राम-आधारित इनिशियलाइजेशन का लाभ उठाकर $\alpha$ -entmax नॉर्मलाइज़र $\tau$ की गणना की लागत को नाटकीय रूप से कम करता है।

मुख्य नवाचार (Core Innovations)

ऑन-चिप हिस्टोग्राम इनिशियलाइजेशन:
- $\tau$ खोजने के लिए कच्चे स्कोर्स (raw scores) पर पुनरावृत्ति करने के बजाय, एल्गोरिदम अटेंशन स्कोर्स को GPU के ऑन-चिप SRAM के माध्यम से स्ट्रीम करता है।
- यह SRAM में पूरी तरह से स्कोर्स का एक कॉम्पैक्ट हिस्टोग्राम (B बिन्स में डिस्क्रीटाइज किया गया) बनाता है। यह हिस्टोग्राम सघन इंटरमीडिएट मैट्रिसेस को भौतिक रूप से बनाए बिना स्कोर्स के वितरण का प्रतिनिधित्व करता है।
- सैद्धांतिक गारंटी: लेखक सिद्ध करते हैं कि इस हिस्टोग्राम पर $\tau$ को हल करने से वास्तविक नॉर्मलाइज़र ( $\tau^*$ ) पर एक प्रमाणित निचली सीमा ( $\tau_h \leq \tau^*$ ) प्राप्त होती है। यह एक उच्च-गुणवत्ता वाला प्रारंभिक अनुमान प्रदान करता है।
तीव्र अभिसरण (1–2 इटरेशन):
- क्योंकि हिस्टोग्राम इनिशियलाइजेशन वास्तविक रूट के अत्यंत निकट होता है, इसलिए एक आगामी हाइब्रिड सॉल्वर (जो $\alpha$ के आधार पर Halley, Newton, और Secant विधियों को जोड़ता है) आमतौर पर 1 से 2 इटरेशन में सटीक समाधान पर अभिसरण करता है।
- यह पिछले तरीकों के विपरीत है जिन्हें कई इटरेशन की आवश्यकता होती थी, जो प्रभावी रूप से एक मल्टी-पास समस्या को लगभग सिंगल-पास ऑपरेशन में बदल देता है।
स्पैरसिटी-अवेयर GPU कार्यान्वयन (Sparsity-Aware GPU Implementation):
- बिट-पैक्ड ब्लॉक मास्किंग: रिफाइनमेंट चरण के दौरान, एल्गोरिदम एक बाइनरी मास्क बनाता है जो उन ब्लॉक्स को इंगित करता है जिनमें नॉन-जीरो वेट्स हैं। इस मास्क को बिट-पैकिंग (32 ब्लॉक्स प्रति int32) का उपयोग करके संग्रहीत किया जाता है, जिससे मेमोरी ओवरहेड न्यूनतम हो जाता है।
- कुशल ट्रैवर्सल (Efficient Traversal): बैकवर्ड पास और अंतिम आउटपुट गणना केवल नॉन-जीरो ब्लॉक्स को ट्रैवर्स करने के लिए नेटिव GPU निर्देशों (fns - फाइंड नेक्स्ट सेट, popc - पॉपुलेशन काउंट) का उपयोग करती है, जिससे ज़ीरो ब्लॉक्स को नगण्य ओवरहेड के साथ छोड़ा जा सकता है।
हाइब्रिड सॉल्वर रणनीति:
- पेपर $\alpha$ $α$ -entmax की स्मूथनेस प्रॉपर्टीज का उपयोग करके इष्टतम रूट-फाइंडिंग स्टेप का चयन करता है:
  - $\alpha \leq 1.5$ : Halley's method का उपयोग करता है (सेकंड डेरिवेटिव का उपयोग करता है)।
  - $1.5 < \alpha \leq 2$ : Newton's method का उपयोग करता है (फर्स्ट डेरिवेटिव का उपयोग करता है)।
  - $\alpha > 2$ : Secant method का उपयोग करता है।
- स्थिरता सुनिश्चित करने के लिए, यदि कोई अपडेट स्टेप वैध सीमाओं से बाहर गिर जाता है, तो एक सुरक्षित बाइसेक्शन (bisection) फॉलबैक का उपयोग किया जाता है।

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

एल्गोरिद्मिक दक्षता: एक हिस्टोग्राम-आधारित इनिशियलाइजेशन पेश किया गया जो $\tau$ गणना के लिए इटरेशन की संख्या को कई से घटाकर आमतौर पर 1–2 कर देता है, जिससे $\alpha$ -entmax को गति में सॉफ्टमैक्स के साथ प्रतिस्पर्धा करने में सक्षम बनाया जा सका।
हार्डवेयर ऑप्टिमाइज़ेशन: एक ऐसा Triton कर्नेल डिज़ाइन किया गया जो GPU मेमोरी पदानुक्रम (HBM $\to$ SRAM) और बिट-पैक्ड डेटा स्ट्रक्चर का लाभ उठाता है ताकि डायनेमिक स्पैरसिटी को कुशलतापूर्वक संभाला जा सके।
सैद्धांतिक गारंटी: हिस्टोग्राम एप्रोक्सिमेशन के लोअर-बाउंड प्रॉपर्टी और विभिन्न $\alpha$ मानों के लिए डेरिवेटिव की निरंतरता (continuity) के लिए प्रमाण प्रदान किए गए, जो हाइब्रिड सॉल्वर की स्थिरता सुनिश्चित करते हैं।
परफॉरमेंस ब्रेकथ्रू: प्रदर्शित किया गया कि जब ब्लॉक स्पैरसिटी मध्यम-से-उच्च (>60%) होती है (जो लॉन्ग-कॉन्टेक्स्ट ट्रेनिंग में एक सामान्य स्थिति है), तो ADASPLASH-2 प्रशिक्षण थ्रूपुट में FlashAttention-2 के बराबर या उससे बेहतर प्रदर्शन करता है।

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

लेखकों ने सिंथेटिक बेंचमार्क, लैंग्वेज मॉडलिंग टास्क और लॉन्ग-कॉन्टेक्स्ट इवैल्यूएशन पर ADASPLASH-2 का मूल्यांकन किया।

रनटाइम दक्षता:
- मध्यम-से-उच्च स्पैरसिटी: उन रेजिम्स में जहाँ >60% ब्लॉक स्पैरसिटी है, ADASPLASH-2 दोनों CUDA और Triton कार्यान्वयन वाले FlashAttention-2 को पीछे छोड़ देता है।
- स्पीडअप: उच्च स्पैरसिटी पर, यह FlashAttention-2 की तुलना में >2x स्पीडअप प्राप्त करता है।
- बैकवर्ड पास: बैकवर्ड पास में सबसे महत्वपूर्ण लाभ देखा गया क्योंकि स्पैरसिटी-अवेयर ट्रैवर्सल ज़ीरो ब्लॉक्स को छोड़ देता है, जिससे फॉरवर्ड पास के थोड़े उच्च नॉर्मलाइजेशन लागत की भरपाई हो जाती है।
लैंग्वेज मॉडलिंग (LLaMA-3 आधारित मॉडल):
- लॉन्ग-कॉन्टेक्स्ट (RULER और HELMET): $\alpha$ -entmax (विशेष रूप से NAPE पोजीशनल एनकोडिंग के साथ) के साथ प्रशिक्षित मॉडल लॉन्ग-कॉन्टेक्स्ट कार्यों पर सॉफ्टमैक्स बेसलाइन की तुलना में काफी बेहतर प्रदर्शन करते हैं (32K टोकन तक)। वे वेरिएबल ट्रैकिंग और वर्ड एक्सट्रैक्शन जैसे कार्यों में पर्याप्त लाभ दिखाते हैं, जिनमें विशिष्ट स्पैन पर सटीक ध्यान देने की आवश्यकता होती है।
- शॉर्ट-कॉन्टेक्स्ट: ADASPLASH-2 के साथ प्रशिक्षित मॉडल मानक शॉर्ट-कॉन्टेक्स्ट बेंचमार्क (जैसे ARC, CSQA, LAMBADA) पर सॉफ्टमैक्स बेसलाइन के बराबर या थोड़ा बेहतर प्रदर्शन करते हैं, जो यह सिद्ध करता है कि स्पैरसिटी डेंस रेजिम्स में प्रदर्शन को कम नहीं करती है।
- इन-कॉन्टेक्स्ट लर्निंग (ICL): जैसे-जैसे कॉन्टेक्स्ट लेंथ बढ़ती है, $\alpha$ -entmax + NAPE बेहतर फ्यू-शॉट प्रदर्शन दिखाता है, जो सुझाव देता है कि स्पैरसिटी मॉडल को प्रासंगिक इन-कॉन्टेक्स्ट उदाहरणों पर ध्यान केंद्रित करने में मदद करती है।

5. महत्व और प्रभाव (Significance and Impact)

स्केलेबिलिटी: ADASPLASH-2 बड़े पैमाने पर स्पारस अटेंशन मैकेनिज्म को प्रशिक्षित करने में एक बड़ी बाधा को दूर करता है। यह शोधकर्ताओं को प्रशिक्षण की गति से समझौता किए बिना $\alpha$ -entmax (डायनेमिक स्पैरसिटी, बेहतर लॉन्ग-कॉन्टेक्स्ट जनरलाइजेशन) के सैद्धांतिक लाभों का उपयोग करने की अनुमति देता है।
लॉन्ग-कॉन्टेक्स्ट समाधान: स्पारस अटेंशन के कुशल प्रशिक्षण को सक्षम करके, यह कार्य बड़े भाषा मॉडलों में "अटेंशन डिस्पर्शन" (attention dispersion) समस्या को हल करने के लिए एक व्यावहारिक मार्ग प्रदान करता है, जहाँ सॉफ्टमैक्स बहुत लंबे अनुक्रमों में प्रासंगिक टोकनों को पहचानने में विफल रहता है।
हार्डवेयर को-डिज़ाइन: यह पेपर ऐसे अटेंशन कर्नेल्स को डिज़ाइन करने के लिए एक ब्लूप्रिंट के रूप में कार्य करता है जो न केवल एल्गोरिद्मिक रूप से सुदृढ़ हैं बल्कि आधुनिक GPU मेमोरी पदानिध (SRAM उपयोग, बिट-पैकिंग और इंस्ट्रक्शन-लेवल पैरेललिज्म) के लिए विशेष रूप से अनुकूलित हैं।

संक्षेप में, ADASPLASH-2 स्पार, डिफरेंशिएबल अटेंशन के सैद्धांतिक लाभों और हाई-परफॉरमेंस GPU ट्रेनिंग की व्यावहारिक आवश्यकताओं के बीच के अंतर को सफलतापूर्वक पाटता है, जिससे यह लॉन्ग-कॉन्टेक्स्ट ट्रांसफॉर्मर ट्रेनिंग के लिए सॉफ्टमैक्स का एक व्यवहार्य विकल्प बन जाता है।