PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training

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

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

मैट्रिक्स फलन (matrix functions), जैसे कि वर्गमूल (square roots), व्युत्क्रम मूल (inverse roots), और ऑर्थोगोनलाइजेशन (ध्रुवीय अपघटन/polar decomposition), न्यूरल नेटवर्क प्रशिक्षण (जैसे कि Shampoo, Muon) में प्रीकंडीशनड ग्रेडिएंट विधियों के लिए मौलिक हैं। जबकि SVD (Singular Value Decomposition) या आइगेनवैल्यू अनुमान की आवश्यकता वाले तरीकों की तुलना में जनरल मैट्रिक्स मल्टीप्लिकेशन (GEMMs) पर आधारित पुनरावृत्ति एल्गोरिदम (iterative algorithms) GPU त्वरण (acceleration) के लिए पसंद किए जाते हैं, मौजूदा अत्याधुनिक दृष्टिकोण महत्वपूर्ण सीमाओं से ग्रस्त हैं:

1. अनुकूलन क्षमता का अभाव (Lack of Adaptivity): हालिया विधियाँ (जैसे कि PolarExpress, CANS) अक्सर एक पूर्व-निर्धारित स्पेक्ट्रल अंतराल (जैसे $[\ell, u]$) के लिए अनुकूलित निश्चित बहुपद अपडेट (fixed polynomial updates) पर निर्भर करती हैं। यदि इनपुट मैट्रिक्स का वास्तविक सिंगुलर वैल्यू वितरण इस अंतराल से विचलित होता है, तो अभिसरण (convergence) नाटकीय रूप से गिर जाता है।
2. पूर्व ज्ञान पर निर्भरता (Dependency on Prior Knowledge): ये विधियाँ आमतौर पर सिंगुलर वैल्यू वितरण के पूर्व ज्ञान या सबसे छोटे और सबसे बड़े सिंगुलर मानों पर कड़े बंधों (tight bounds) की आवश्यकता रखती हैं, जो अक्सर उपलब्ध नहीं होते या स्वयं मैट्रिक्स फलन का अनुमान लगाने जितने ही महंगे होते हैं।
3. सीमित दायरा (Narrow Scope): कई हालिया त्वरण तकनीकें विशिष्ट उपयोग के मामलों (जैसे कि केवल Muon ऑप्टिमाइज़र में ध्रुवीय अपघटन) के लिए तैयार की गई हैं, न कि विभिन्न मैट्रिक्स फलनों के लिए एक सामान्य ढांचे के रूप में।

कार्यप्रणाली: PRISM ढांचा (Methodology: The PRISM Framework)

लेखक PRISM (Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation) का प्रस्ताव करते हैं, जो एक सामान्य मेटा-एल्गोरिदम है जिसे बिना किसी स्पष्ट स्पेक्ट्रल बाउंड की आवश्यकता के, इट्रेटिव मैट्रिक्स फंक्शन गणनाओं को तेज करने के लिए डिज़ाइन किया गया है। PRISM दो अलग-अलग भागों में कार्य करता है:

भाग I: बुनियादी सेटअप (पुनरावृत्ति बहुपद सन्निकटन - Iterative Polynomial Approximation)

PRISM मैट्रिक्स फलन $T(A)$ की गणना को एक इट्रेटिव बहुपद सन्निकटन समस्या के रूप में फ्रेम करता है।

• $T(a)$ को अनुमानित करने वाले स्केलर $x$ के लिए, समस्या को एक अवशिष्ट फलन (residual function) $f(\xi)$ को अनुमानित करने के रूप में पुनर्गठित किया जाता है, जहाँ $\xi$ समाधान से दूरी को मापता है।
• मानक इट्रेटिव विधियाँ (जैसे Newton-Schulz) $\xi=0$ के आसपास एक निश्चित-क्रम टेलर बहुपद (fixed-order Taylor polynomial) $f_d(\xi)$ का उपयोग करके प्राप्त की जाती हैं।
• मैट्रिक्स इटरेशन का रूप $X_{k+1} = X_k f_d(R_k)$ होता है, जहाँ $R_k$ एक अवशिष्ट मैट्रिक्स (residual matrix) है (जैसे कि साइन फंक्शन के लिए $I - X_k^2$)।

भाग II: त्वरण (अनुकूली फिटिंग और स्केचिंग - Adaptive Fitting and Sketching)

जब $R_k$ का स्पेक्ट्रम 0 के करीब नहीं होता है, तो निश्चित टेलर बहुपदों की धीमी प्रारंभिक अभिसरण दर को दूर करने के लिए, PRISM दो तंत्र पेश करता है:

1. अनुकूली बहुपद फिटिंग (Adaptive Polynomial Fitting):
   निश्चित टेलर बहुपद के बजाय, PRISM उच्चतम-क्रम के गुणांक को एक सीखने योग्य पैरामीटर $\alpha$ द्वारा बदल देता है। वर्तमान $R_k$ के स्पेक्ट्रम में फिट होने के लिए प्रत्येक इटरेशन में बहुपद को अपडेट किया जाता है।

• अपडेट बहुपद को $g_d(\xi; \alpha) = f_{d-1}(\xi) + \alpha \xi^d$ के रूप में परिभाषित किया गया है।
• इष्टतम $\alpha_k^*$ को $\|I - X_{k+1}^2\|_F^2$ (या अन्य फलनों के लिए समकक्ष) के अवशिष्ट नॉर्म को न्यूनतम करके पाया जाता है। यह प्रभावी रूप से $R_k$ के आइगेनवैल्यूज को लीस्ट-स्क्वायर्स सेंस में फिट करता है।
• स्थिरता सुनिश्चित करने के लिए, $\alpha_k$ को एक निश्चित अंतराल $[\ell, u]$ (जैसे $d=1$ के लिए $[1/2, 1]$) में सीमित किया गया है, जो इनपुट स्पेक्ट्रम से स्वतंत्र अभिसरण गुणों की गारंटी देता है।

2. यादृच्छिक स्केचिंग (Randomized Sketching):
   सटीक इष्टतम $\alpha_k^*$ की गणना करने के लिए अवशिष्ट मैट्रिक्स की उच्च घातों के ट्रेस (traces) का मूल्यांकन करना आवश्यक है, जो कि कम्प्यूटेशनल रूप से महंगा ($O(n^3)$ या उससे अधिक) है।

• PRISM एक स्केच किए गए संस्करण पर लीस्ट-स्क्वायर्स समस्या को हल करने के लिए यादृच्छिक स्केचिंग (Oblivious Subspace Embeddings, जैसे कि Gaussian matrices का उपयोग करके) का उपयोग करता है।
• स्केच किए गए अवशिष्ट पर मिनिमाइजेशन समस्या को हल करके, $\alpha_k$ खोजने की लागत को $O(n^2 \log n)$ (या $O(n^2 p)$ जहाँ $p \ll n$) तक कम कर दिया जाता है, जो मुख्य इटरेशन में मैट्रिक्स गुणन की $O(n^3)$ लागत की तुलना में नगण्य है।
• सैद्धांतिक गारंटी (Theorem 2) दिखाती है कि $p = O(\log n)$ के स्केच आयाम के साथ, अभिसरण दर अनिवार्य रूप से सटीक मिनिमाइजेशन के समान रहती है।

प्रमुख योगदान (Key Contributions)

1. सामान्य त्वरण ढांचा (General Acceleration Framework): PRISM एक एकीकृत दृष्टिकोण प्रदान करता है जो Newton-Schulz (वर्गमूल, व्युत्क्रम मूल, साइन और ध्रुवीय अपघटन के लिए), चेबिशेव विधियों (Chebyshev methods), और इनवर्स न्यूटन विधियों सहित व्यापक श्रेणी के इट्रेटिव एल्गोरिदम पर लागू होता है।
2. डिस्ट्रीब्यूशन-फ्री एडेप्टिविटी (Distribution-Free Adaptivity): पूर्व कार्य के विपरीत, PRISM पूर्व सिंगुलर वैल्यू बाउंड्स या डिस्ट्रीब्यूशन के ज्ञान की आवश्यकता के बिना, वर्तमान इटरेट के विकसित होते स्पेक्ट्रम के अनुकूल गतिशील रूप से अनुकूलित होता है।
3. कुशल कार्यान्वयन (Efficient Implementation): रैंडमाइज्ड स्केचिंग का एकीकरण यह सुनिश्चित करता है कि एडेप्टिव बहुपद फिटिंग का ओवरहेड न्यूनतम हो, जिससे GEMM-प्रधान इटरेशन की GPU-फ्रेंडली प्रकृति बनी रहे।
4. सैद्धांतिक गारंटी (Theoretical Guarantees): पेपर अभिसरण प्रमाण प्रदान करता है जो दिखाते हैं कि PRISM-त्वरित इटरेशन पर्याप्त बड़े स्केच आयाम के साथ भी द्विघातीय (quadratic) या उच्च-क्रम अभिसरण दर बनाए रखते हैं।

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

लेखक सिंथेटिक बेंचमार्क और न्यूरल नेटवर्क ऑप्टिमाइज़र में एकीकरण के माध्यम से PRISM को सत्यापित करते हैं:

• अभिसरण गति (Convergence Speed): विभिन्न आस्पेक्ट रेश्यो और स्पेक्ट्रल डिस्ट्रीब्यूशन वाली मैट्रिसेस (Gaussian random matrices और Heavy-Tailed Marchenko-Pastur distributions सहित जो प्री-ट्रेन्ड मॉडल्स की नकल करते हैं) पर, PRISM लगातार क्लासिकल Newton-Schulz और निश्चित-अंतराल विधियों जैसे PolarExpress से बेहतर प्रदर्शन करता है।
  • PolarExpress अक्सर धीमा हो जाता है या अभिसरण करने में विफल रहता है जब वास्तविक सिंगुलर वैल्यू रेंज इसके अनुकूलित अंतराल से विचलित होती है।
  • PRISM पूरे रेंज में स्थिर स्पीडअप बनाए रखता है, $\alpha_k$ को स्वचालित रूप से अनुकूलित करता है।
• न्यूरल नेटवर्क प्रशिक्षण (Neural Network Training):
  • Shampoo Optimizer: CIFAR डेटासेट पर ResNet-20/32 के लिए इनवर्स स्क्वायर रूट प्रीकंडीशनर्स को कंप्यूट करने के लिए PRISM को एकीकृत करने से आइगेन-डिकम्पोजिशन और PolarExpress बेसलाइन्स की तुलना में तेज़ प्रशिक्षण समय प्राप्त हुआ।
  • Muon Optimizer: GPT-2 ट्रेनिंग (FineWeb डेटासेट) में ग्रेडिएंट मैट्रिसेस के ध्रुवीय अपघटन पर PRISM को लागू करने से AdamW, PolarExpress और फिक्स्ड-डिग्री Newton-Schulz वेरिएंट की तुलना में बेहतर वैलिडेशन लॉस और प्रशिक्षण दक्षता प्रदर्शित हुई।

महत्व और दावे (Significance and Claims)

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