तकनीकी सारांश: RoPE अटेंशन को लगभग रैखिक समय में प्रशिक्षित किया जा सकता है

समस्या की परिभाषा

रोटरी पोजीशन एम्बेडिंग (RoPE) तंत्र अत्याधुनिक लार्ज लैंग्वेज मॉडल्स (LLMs) जैसे कि Llama, Claude और Apple के मॉडल्स में एक मानक घटक बन गया है, जो पारंपरिक पोजीशनल एनकोडिंग की तुलना में टोकन संबंधों को पकड़ने में बेहतर अभिव्यक्ति प्रदान करता है। हालाँकि, RoPE में निहित स्थिति-निर्भर रोटेशन (position-dependent rotations) अटेंशन मैकेनिज्म की गणना को जटिल बना देते हैं।

जबकि हाल के कार्य ([AS24a]) ने "बाउंडेड एंट्री" (bounded entry) शासन के तहत (जहाँ मैट्रिक्स प्रविष्टियाँ एक पैरामीटर $B$ द्वारा सीमित होती हैं) RoPE अटेंशन के फॉरवर्ड (forward) कंप्यूटेशन के लिए एक लगभग रैखिक समय एल्गोरिदम ($n^{1+o(1)}$) स्थापित किया था, बैकवर्ड (backward) कंप्यूटेशन (प्रशिक्षण के लिए ग्रेडिएंट गणना) अनसुलझा रहा। बैकवर्ड कंप्यूटेशन स्वाभाविक रूप से अधिक जटिल है क्योंकि इसमें अटेंशन मैट्रिक्स और पोजीशनल एम्बेडिंग्स का गैर-रैखिक रूपांतरण शामिल होता है। यह कार्य जिस केंद्रीय प्रश्न को संबोधित करता है वह यह है कि क्या बाउंडेड एंट्री स्थितियों के तहत RoPE अटेंशन के लिए बैकवर्ड ग्रेडिएंट कंप्यूटेशन, फॉरवर्ड कंप्यूटेशन की तरह ही लगभग रैखिक समय दक्षता प्राप्त कर सकता है।

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

लेखक बैकवर्ड RoPE अटेंशन कंप्यूटेशन के लिए पहला एल्गोरिदम विकसित करते हैं जो लगभग रैखिक समय में चलता है। यह दृष्टिकोण क्लोज्ड-फॉर्म ग्रेडिएंट व्युत्पत्ति (closed-form gradient derivation), लो-रैंक एप्रोक्सिमेशन (low-rank approximation), बहुपद विधियों (polynomial methods) और फास्ट फूरियर ट्रांसफॉर्म (FFT) के संयोजन पर निर्भर करता है।

1. क्लोज्ड-फॉर्म ग्रेडिएंट रिफॉर्मुलेशन

पेपर सबसे पहले वेट मैट्रिसेस के सापेक्ष RoPE अटेंशन लॉस फंक्शन के ग्रेडिएंट के लिए एक क्लोज्ड-फॉर्म अभिव्यक्ति व्युत्पन्न करता है। "टेन्सर ट्रिक" (क्रोनिकर उत्पाद) का उपयोग करते हुए और अटेंशन मैट्रिक्स $A(X)$ को रिफॉर्मुलेट करते हुए, ग्रेडिएंट को इस प्रकार व्यक्त किया जाता है:
$$\frac{d\text{Loss}(x)}{dx} = \tilde{A}^\top \text{vec}(\gamma(x))$$
जहाँ $\gamma(x)$ एक जटिल मैट्रिक्स फंक्शन है जिसमें शामिल हैं:

• $s(x)$: सामान्यीकृत सॉफ्टमैक्स (Softmax) वेक्टर।
• $\ell(x)$: अटेंशन आउटपुट और टारगेट के बीच के अंतर से प्राप्त त्रुटि पद (error term)।
• $\beta(x)$: त्रुटि और वैल्यू मैट्रिक्स को संयोजित करने वाला एक पद।
• $\gamma(x)$: $s(x)$ के डायगोनल और $s(x)s(x)^\top$ के बाहरी उत्पाद (outer product) को शामिल करने वाला एक पद जो $\beta(x)$ पर कार्य करता है।

2. लो-रैंक एप्रोक्सिमेशन रणनीति

लगभग रैखिक समय जटिलता प्राप्त करने के लिए, लेखक $\gamma(x)$ के घटकों को लो-रैंक मैट्रिसेस का उपयोग करके अनुमानित (approximate) करते हैं। रणनीति $\gamma(x)$ को दो भागों, $\gamma_1(x)$ और $\gamma_2(x)$ में विभाजित करने और प्रत्येक को अलग से अनुमानित करने की है:

• $s(x)$ और $\ell(x)$ का अनुमान लगाना: [AS24a] से फॉरवर्ड एल्गोरिदम पर निर्माण करते हुए, लेखक दिखाते हैं कि सामान्यीकृत सॉफ्टमैक्स $s(x)$ को $n^{1+o(1)}$ समय में लो-रैंक मैट्रिसेस $U_1 V_1^\top$ द्वारा अनुमानित किया जा सकता है। इसके बाद त्रुटि पद $\ell(x)$ को इस परिणाम का उपयोग करके अनुमानित किया जाता है।
• $\beta(x)$ का अनुमान लगाना: चूंकि $\beta(x)$ अपने घटकों और वैल्यू मैट्रिक्स के उत्पाद से बना है, इसलिए इसे इसके घटकों के अनुमानों के आधार पर लो-रैंक फैक्टर्स का निर्माण करके अनुमानित किया जाता है।
• $\gamma(x)$ का अनुमान लगाना:
  • $\gamma_1(x) = \text{diag}(s(x))\beta(x)$ को $s(x)$ और $\beta(x)$ के लो-रैंक फैक्टर्स को रो-वाइज क्रोनिकर उत्पादों का उपयोग करके संयोजित करके अनुमानित किया जाता है।
  • $\gamma_2(x) = s(x)s(x)^\top\beta(x)$ को मध्यवर्ती पदों (intermediate terms) को प्रीकंप्यूट करके और $s(x)$ तथा $\beta(x)$ की लो-रैंक संरचना का उपयोग करके अनुमानित किया जाता है।

3. हार्डनेस एनालिसिस (Hardness Analysis)

बाउंडेड एंट्री स्थिति की आवश्यकता को स्थापित करने के लिए, लेखक स्ट्रॉन्ग एक्सपोनेंशियल टाइम हाइपोथेसिस (SETH) पर आधारित लोअर बाउंड्स व्युत्पन्न करते हैं। वे सिद्ध करते हैं कि यदि एंट्री बाउंड $B$ एक निश्चित सीमा (विशेष रूप से $B = \omega(\sqrt{\log n})$) से अधिक है, तो SETH मानते हुए कोई भी एल्गोरिदम सब-क्वाड्रेटिक समय ($O(n^{2-q})$) में ग्रेडिएंट की गणना नहीं कर सकता है। यह पुष्टि करता है कि बाउंडेड एंट्री धारणा केवल एक तकनीकी सुविधा नहीं है, बल्कि सब-क्वाड्रेटिक प्रदर्शन के लिए एक मौलिक आवश्यकता है।

मुख्य योगदान

1. क्लोज्ड-फॉर्म ग्रेडिएंट: पेपर RoPE अटेंशन के ग्रेडिएंट के लिए पहला क्लोज्ड-फॉर्म फॉर्मुलेशन (लेम्मा 4.1) प्रदान करता है और इसकी सटीक समय जटिलता का विश्लेषण करता है, जिससे नैव (naive) गणना में क्वाड्रेटिक बॉटलनेक की पहचान होती है।
2. लगभग रैखिक समय एल्गोरिदम: लेखक बाउंडेड एंट्री स्थितियों के तहत $n^{1+o(1)}$ समय में RoPE अटेंशन के बैकवर्ड ग्रेडिएंट को अनुमानित करने वाले पहले एल्गोरिदम को प्रस्तुत करते हैं (थ्योरम 5.7)। यह फॉरवर्ड पास की दक्षता से मेल खाता है।
3. सैद्धांतिक लोअर बाउंड्स: यह कार्य स्थापित करता है कि सब-क्वाड्रेटिक प्रदर्शन के लिए बाउंडेड एंट्री की स्थिति आवश्यक है, जो SETH से प्राप्त हार्डनेस रिजल्ट (थ्योरम 6.1) प्रदान करती है।
4. एल्गोरिद्मिक तकनीकें: दृष्टिकोण विशेष रूप से RoPE के संरचनात्मक बाधाओं के लिए तैयार की गई लो-रैंक एप्रोक्सिमेशन तकनीकों के साथ बहुपद अनुमान विधियों (polynomial approximation methods) और FFT को एकीकृत करता है।

परिणाम

मुख्य परिणाम (थ्योरम 5.7) प्रदर्शित करता है कि $d = O(\log n)$ और $B = o(\sqrt{\log n})$ मापदंडों के लिए, $n^{1+o(1)}$ समय में $1/\text{poly}(n)$ द्वारा सीमित एडिटिव एरर (additive error) के साथ RoPE अटेंशन ग्रेडिएंट कंप्यूटेशन समस्या को हल करने के लिए एक एल्गोरिदम मौजूद है।

इसके विपरीत, हार्डनेस रिजल्ट (थ्योरम 6.1) दिखाता है कि यदि $B = \omega(\sqrt{\log n})$ है, तो SETH धारणा के तहत $O(n^{2-q})$ समय में ग्रेडिएंट की गणना करना असंभव है।

महत्व

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

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