💬 NLP

On Fine-Grained I/O Complexity of Attention Backward Passes

यह शोध पत्र रेड-ब्लू पेबल गेम फ्रेमवर्क का उपयोग करके सभी कैश आकारों के लिए अटेंशन बैकवर्ड पासेज के लिए सटीक I/O जटिलता सीमाएं स्थापित करता है, बड़े-कैश परिदृश्यों में फ्लैशअटेंशन (FlashAttention) की इष्टतमता को मान्य करता है, और एक नवीन एल्गोरिदम प्रस्तावित करता है जो छोटे-कैश वातावरण के लिए सैद्धांतिक इष्टतमता प्राप्त करता है और इन परिणामों को स्पार्स अटेंशन (sparse attention) तक विस्तारित करता है।

मूल लेखक: Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Song Yue, Jiahao Zhang

प्रकाशित 2026-01-26

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

CC BY 4.0

मूल लेखक: Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Song Yue, Jiahao Zhang

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

कल्पना कीजिए कि आप एक मास्टर शेफ (AI मॉडल) हैं जो मेहमानों की एक बहुत लंबी सूची (शब्दों का "कॉन्टेक्स्ट" या अनुक्रम) के लिए एक विशाल भोज बनाने की कोशिश कर रहे हैं। अपने व्यंजन को उत्तम बनाने के लिए, आपको यह तय करने के लिए कि प्रत्येक सामग्री का कितना उपयोग करना है, हर मेहमान की पसंद की तुलना दूसरे हर मेहमान की पसंद से करनी होगी। यह लार्ज लैंग्वेज मॉडल्स में "अटेंशन" (Attention) मैकेनिज्म है।

समस्या क्या है? जैसे-जैसे मेहमानों की सूची बढ़ती है, आवश्यक जाँचों की संख्या विस्फोट की तरह बढ़ती जाती है। यदि आपके पास 1,000 मेहमान हैं, तो आप दस लाख जाँच करेंगे। यदि आपके पास 10,000 हैं, तो आप 10 करोड़ जाँच करेंगे। यह वह "क्वाड्रेटिक स्केलिंग" (quadratic scaling) बाधा है जिसका उल्लेख पेपर में किया गया है।

अब, कल्पना कीजिए कि आपके किचन में दो प्रकार के स्टोरेज हैं:

काउंटरटॉप (कैश): चूल्हे के ठीक बगल में एक छोटा, तेज़ और महंगा स्थान जहाँ से आप तुरंत सामग्री उठा सकते हैं।
पेंट्री (मेमोरी): एक विशाल, धीमी और गहरी स्टोरेज रूम जहाँ आपकी सभी सामग्रियाँ रखी जाती हैं।

हर बार जब आपको पेंट्री से काउंटरटॉप तक सामग्री लेने के लिए जाना पड़ता है, तो इसमें आपका समय और ऊर्जा खर्च होती है। यह पेंट्री के चक्कर लगाने को कंप्यूटर वैज्ञानिक I/O कॉम्प्लेक्सिटी (इनपुट/आउटपुट) कहते हैं। लक्ष्य इन चक्करों को कम करना है।

मुख्य समस्या: "बैकवर्ड पास" (The Backward Pass)

जब शेफ सीख रहा होता है (ट्रेनिंग), तो वे केवल व्यंजन नहीं बनाते; उन्हें यह भी समझना होता है कि क्या गलत हुआ ताकि वे अगली बार के लिए रेसिपी को सुधार सकें। इसे बैकवर्ड पास कहा जाता है।

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

खोज: काउंटरटॉप के आकार पर निर्भर करता है

लेखकों ने महसूस किया कि उत्तर पूरी तरह से इस बात पर निर्भर करता है कि आपके रेसिपी (हिडन डायमेंशन, $d$ ) के आकार की तुलना में आपका काउंटरटॉप (कैश) कितना बड़ा है। उन्होंने $d^2$ के एक विशिष्ट आकार पर एक "टिपिंग पॉइंट" (बदलाव का बिंदु) पाया।

1. "बड़ा काउंटरटॉप" परिदृश्य ( $M \ge d^2$ )

यदि आपका काउंटरटॉप इतना बड़ा है कि वह एक बार में आपकी सामग्रियों का एक महत्वपूर्ण हिस्सा रख सके, तो फ्लैशअटेंशन एकदम सही है।

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

2. "छोटा काउंटरटॉप" परिदृश्य ( $M < d^2$ )

यदि आपका काउंटरटॉप बहुत छोटा है (जैसे पुराने या सस्ते कंप्यूटरों पर), तो फ्लैशअटेंशन लड़खड़ाने लगता है। यह एक ऐसी रणनीति का उपयोग करने की कोशिश करता है जो बड़े काउंटरों के लिए काम करती है, जिससे इसे पेंट्री के अनावश्यक चक्कर लगाने पड़ते हैं।

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

"स्पार्स" (Sparse) ट्विस्ट

पेपर ने स्पार्स अटेंशन (Sparse Attention) नामक एक भिन्नता को भी देखा।

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

पेपर के दावों का सारांश

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

संक्षेप में, पेपर कहता है: "यदि आपके पास बड़ा किचन है, तो फ्लैशअटेंशन के साथ बने रहें। यदि आपके पास छोटा किचन है, तो समय और ऊर्जा बचाने के लिए हमारी नई विधि पर स्विच करें।"

तकनीकी सारांश: अटेंशन बैकवर्ड पासेस की फाइन-ग्रेन्ड I/O जटिलता पर

समस्या विवरण

लार्ज लैंग्वेज मॉडल्स (LLMs) अत्यधिक रूप से सेल्फ-अटेंशन मैकेनिज्म पर निर्भर करते हैं, जो अनुक्रम लंबाई $n$ के सापेक्ष द्विघाती कम्प्यूटेशनल जटिलता ( $O(n^2)$ ) प्रदर्शित करता है। जबकि फ्लैशअटेंशन (FlashAttention) ने I/O जागरूकता (फास्ट ऑन-चिप कैश और स्लो हाई-बैंडविड्थ मेमोरी के बीच डेटा ट्रांसफर को कम करके) का लाभ उठाकर फॉरवर्ड पास के अनुकूलन के लिए उद्योग मानक बन गया है, बैकवर्ड पास के लिए सैद्धांतिक समझ अभी भी अपूर्ण है।

मौजूदा साहित्य बड़े कैश आकार ( $M \in [d, nd]$ ) के लिए फ्लैशअटेंशन के बैकवर्ड पास की I/O जटिलता के ऊपरी बाउंड्स (upper bounds) प्रदान करता है, लेकिन मिलान करने वाले निचले बाउंड्स (lower bounds) का अभाव है। फलस्वरूप, यह अज्ञात है कि क्या फ्लैशअटेंशन सभी कैश आकारों के स्पेक्ट्रम में, विशेष रूप से छोटे-कैश परिदृश्यों ( $M < d^2$ ) या स्पार्स अटेंशन मैकेनिज्म के लिए, बैकवर्ड गणनाओं के लिए इष्टतम (optimal) है। यह शोध पत्र इस खुले प्रश्न को संबोधित करता है: किसी भी कैश आकार के लिए अटेंशन बैकवर्ड गणनाओं की इष्टतम I/O जटिलता क्या है?

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

लेखक दो-स्तरीय मेमोरी पदानुक्रम (फास्ट कैश आकार $M$ और स्लो मेमोरी) को मॉडल करने के लिए रेड-ब्लू पेबल गेम [HK81] पर आधारित एक कठोर सैद्धांतिक ढांचे का उपयोग करते हैं। यह ढांचा गणना को एक निर्देशित अचक्रीय ग्राफ (DAG) के रूप में मानता है जहाँ रेड पेबल्स कैश में डेटा का प्रतिनिधित्व करते हैं और ब्लू पेबल्स मेमोरी में डेटा का प्रतिनिधित्व करते हैं। I/O जटिलता को रेड पेबल्स की संख्या (कैश आकार) पर प्रतिबंध दिए बिना आउटपुट को कंप्यूट करने के लिए आवश्यक न्यूनतम इनपुट/आउटपुट ऑपरेशन्स के रूप में परिभाषित किया गया है।

मुख्य कार्यप्रणाली संबंधी चरण शामिल हैं:

कम्प्यूटेशनल ग्राफ विश्लेषण: लेखक अटेंशन बैकवर्ड पास के लिए एक विस्तृत कम्प्यूटेशनल ग्राफ का निर्माण करते हैं, जिसमें अटेंशन वेट्स ( $X$ और $Y$ ) के ग्रेडिएंट कंप्यूटेशन को स्पष्ट रूप से मॉडल किया गया है। इसमें मानक मैट्रिक्स मल्टीप्लिकेशन (विशेष रूप से $A_1 X A_2^\top$ ) में निहित समेशन ट्रीज़ (summation trees) का विश्लेषण करना शामिल है।
मामला विभेदन (Case Differentiation): विश्लेषण कैश आकार $M$ $M$ और हेड डायमेंशन के वर्ग $d^2$ $d^{2}$ के बीच संबंध के आधार पर दो रेजिम में विभाजित होता है:
- छोटा कैश (Small Cache): $M = o(d^2)$
- बड़ा कैश (Large Cache): $M = \Omega(d^2)$
एल्गोरिदम डिज़ाइन: छोटे-कैश रेजिम के लिए, लेखक एक नया एल्गोरिदम (एल्गोरिदम 6) प्रस्तावित करते हैं जो फ्लैशअटेंशन की ब्लॉक-वाइज रणनीति से अलग, $n \times n$ अटेंशन मैट्रिसेस और $n \times d$ इंटरमीडिएट मैट्रिसेस को स्पष्ट रूप से पढ़ता और लिखता है।
स्पार्स अटेंशन विस्तार: फ्रेमवर्क को स्पार्स अटेंशन के लिए विस्तारित किया गया है, जहाँ मैट्रिसेस को नॉन-जीरो एंट्रीज द्वारा स्टोर किया जाता है, जिससे नॉन-जीरो तत्वों ( $Z_{input}$ , $Z_{QK}$ ) के आधार पर लोअर बाउंड्स स्थापित होते हैं।

मुख्य योगदान और परिणाम

1. डेंस अटेंशन के लिए टाइट I/O जटिलता बाउंड्स

यह शोध पत्र सभी कैश आकारों $M$ के लिए अटेंशन बैकवर्ड पासेस की I/O जटिलता के लिए मिलान करने वाले ऊपरी और निचले बाउंड्स स्थापित करता है। मुख्य परिणाम (थ्योरम 1.1) यह है कि I/O जटिलता है:
$\Theta\left(\min\left\{ \frac{n^2d^2 + nd^3}{M}, \frac{n^2d + nd^2}{\sqrt{M}} \right\}\right)$

यह परिणाम $M = \Theta(d^2)$ पर एक महत्वपूर्ण ट्रांजिशन पॉइंट प्रकट करता है:

बड़ा कैश ( $M = \Omega(d^2)$ ): जटिलता $\Theta(\frac{n^2d^2 + nd^3}{M})$ है। लेखक सिद्ध करते हैं कि इस रेजिम में फ्लैशअटेंशन इष्टतम (optimal) है, क्योंकि इसका मौजूदा अपर बाउंड व्युत्पन्न लोअर बाउंड से मेल खाता है।
छोटा कैश ( $M = o(d^2)$ ): जटिलता $\Theta(\frac{n^2d + nd^2}{\sqrt{M}})$ है। लेखक प्रदर्शित करते हैं कि इस रेजिम में फ्लैशअटेंशन इष्टतम नहीं है। वे एक नया एल्गोरिदम (एल्गोरिदम 6) प्रस्तावित करते हैं जो टाइट बाउंड $O(\frac{n^2d + nd^2}{\sqrt{M}})$ प्राप्त करता है, जो फ्लैशअटेंशन के $O(\frac{n^2d^2 + nd^3}{M})$ बाउंड से बेहतर प्रदर्शन करता है।

2. स्पार्स अटेंशन लोअर बाउंड्स

यह शोध पत्र स्पार्स अटेंशन के फॉरवर्ड और बैकवर्ड दोनों पासेस के लिए फाइन-ग्रेन्ड लोअर बाउंड्स प्रदान करता है (थ्योरम 4.5)। इनपुट में $Z_{input}$ नॉन-जीरो एंट्रीज और की-क्वेरी मैट्रिक्स में $Z_{QK}$ नॉन-जीरो एंट्रीज वाली स्पार्स मैट्रिसेस के लिए, I/O जटिलता निम्न द्वारा सीमित है:
$\Omega\left(\min\left\{ \frac{Z_{input}^2}{M}, \frac{Z_{input}\sqrt{Z_{QK}}}{\sqrt{M}} \right\}\right)$
यह डेंस केस परिणामों का सामान्यीकरण करता है और पुष्टि करता है कि छोटे और बड़े कैश व्यवहार के बीच विभाजक बिंदु ( $M = Z_{input}^2 / Z_{QK}$ ) सुसंगत रहता है।

3. सैद्धांतिक पूर्णता

इन बैकवर्ड पास परिणामों को मौजूदा फॉरवर्ड पास विश्लेषण [SY24] के साथ जोड़कर, यह शोध पत्र अटेंशन मैकेनिज्म की I/O जटिलता के सैद्धांतिक चित्र को पूरा करता है। यह पुष्टि करता है कि जहाँ फ्लैशअटेंशन बड़े कैश के लिए इष्टतम है, वहीं सैद्धांतिक इष्टतमता प्राप्त करने के लिए छोटे कैश के लिए अलग-अलग एल्गोरिथमिक रणनीतियों की आवश्यकता होती है।

महत्व

यह शोध पत्र अटेंशन मैकेनिज्म में I/O जटिलता के संबंध में सैद्धांतिक ढांचे को सुदृढ़ करने का दावा करता है। इसका महत्व निम्नलिखित में निहित है:

सिद्धांत को पूरा करना: सभी कैश आकारों में अटेंशन बैकवर्ड पासेस के लिए पहले मिलान करने वाले ऊपरी और निचले बाउंड्स प्रदान करके, मौजूदा साहित्य के अंतराल को दूर करना।
एल्गोरिदम मार्गदर्शन: यह पहचानना कि वर्तमान उद्योग मानक (फ्लैशअटेंशन) छोटे-कैश वातावरण (जैसे पुराने GPU या उच्च-आयामी सेटिंग्स जहाँ $d$ के सापेक्ष $M$ बड़ा है) के लिए उप-इष्टतम (suboptimal) है। प्रस्तावित एल्गोरिदम 6 इन विशिष्ट हार्डवेयर बाधाओं के लिए एक सैद्धांतिक रूप से बेहतर विकल्प प्रदान करता है।
स्पार्स अटेंशन फाउंडेशन: स्पार्स अटेंशन I/O जटिलता के लिए पहले सैद्धांतिक लोअर बाउंड्स प्रदान करना, जो लॉन्ग-कॉन्टेक्स्ट मॉडलिंग में स्पार्सिफिकेशन तकनीकों की सीमाओं को समझने के लिए महत्वपूर्ण है।

लेखक निष्कर्ष निकालते हैं कि उनके निष्कर्ष कुशल LLM ट्रेनिंग और इन्फरेंस सिस्टम विकसित करने के लिए महत्वपूर्ण मार्गदर्शन प्रदान करते हैं, विशेष रूप से जैसे-जैसे हार्डवेयर बाधाएं और मॉडल आयाम विकसित हो रहे हैं।