A GPU-Accelerated JAX Framework for Robust Parametric Component Separation and Clustering Optimization for CMB Polarization Satellites

"CMB ध्रुवीकरण उपग्रहों के लिए सुदृढ़ पैरामीट्रिक घटक पृथक्करण और क्लस्टरिंग अनुकूलन हेतु एक GPU-त्वरित JAX फ्रेमवर्क" का विस्तृत तकनीकी सारांश।

1. समस्या विवरण

कॉस्मिक माइक्रोवेव बैकग्राउंड (CMB) में मौलिक B-मोड ध्रुवीकरण (primordial B-mode polarization) का पता लगाना इन्फ्लेशनरी भौतिकी और टेंसर-टू-स्केलर अनुपात ($r$) की जांच करने के लिए आवश्यक है। हालांकि, यह संकेत अत्यधिक रूप से ध्रुवीकृत खगोलीय अग्रभूमि (astrophysical foregrounds - सिंक्रोट्रॉन और थर्मल डस्ट) द्वारा दूषित होता है, जो स्थानिक रूप से परिवर्तनशील स्पेक्ट्रल एनर्जी डिस्ट्रीबशन्स (SEDs) प्रदर्शित करते हैं।

वर्तमान पैरामीट्रिक घटक पृथक्करण विधियां (जैसे, fgbuster) एक गंभीर बायस-वेरिएंस ट्रेड-ऑफ (bias-variance trade-off) का सामना करती हैं:

• अल्प-पैरामीटरीकरण (Under-parameterization): स्पेक्ट्रल मापदंडों को फिट करने के लिए बड़े, समान आकाश क्षेत्रों का उपयोग करने से सांख्यिकीय शोर (statistical noise) कम होता है, लेकिन यह महत्वपूर्ण व्यवस्थित बायस (systematic bias) उत्पन्न करता है क्योंकि यह स्थानीय अग्रभूमि विविधताओं को पकड़ने में विफल रहता है।
• अति-पैरामीटरीकरण (Over-parameterization): पिक्सेल-दर-पिक्सेल मापदंडों को फिट करना विविधताओं को पकड़ता है, लेकिन यह सांख्यिकीय शोर को बढ़ा देता है, जिससे कॉस्मोलॉजिकल बाधाएं (cosmological constraints) कमजोर हो जाती हैं।

इन प्रभावों को संतुलित करने के लिए इष्टतम स्थानिक विभाजन (patch configuration) खोजने का कार्य पारंपरिक रूप से गणनात्मक रूप से कठिन रहा है, क्योंकि इसमें संभावित आकाश विभाजनों का विशाल खोज स्थान शामिल है, जिसमें अक्सर दस हजार से अधिक CPU घंटे लग जाते हैं।

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

लेखक एक नवीन फ्रेमवर्क प्रस्तावित करते हैं जो FURAX पर आधारित है, जो CMB डेटा विश्लेषण के लिए एक JAX-संचालित वातावरण है, जिसे एंड-टू-एंड डिफरेंशिएबिलिटी (end-to-end differentiability) और GPU त्वरण सक्षम करने के लिए डिज़ाइन किया गया है।

A. पैरामीट्रिक मॉडल और स्थानिक क्लस्टरिंग

• मॉडल: आकाश को घटकों (CMB, डस्ट, सिंक्रोट्रॉन) के एक रैखिक संयोजन के रूप में मॉडल किया गया है जिसमें योगात्मक गाऊसी शोर (additive Gaussian noise) है। मिक्सिंग मैट्रिक्स स्पेक्ट्रल मापदंडों ($\beta_d, T_d, \beta_s$) पर निर्भर करती है।
• पैचिंग रणनीति: निश्चित HEALPix सुपर-पिक्सेल के बजाय, लेखक गोलीय K-मीन्स क्लस्टरिंग (via jax-healpy) का उपयोग करके आकाश को विलगित पैच (disjoint patches) में विभाजित करते हैं। प्रत्येक पैच स्पेक्ट्रल मापदंडों के एक सामान्य सेट को साझा करता है।
• लचीलापन: आकाश को तीन अलग-अलग गैलेक्टिक क्षेत्रों (High, Mid, Low latitude) में विभाजित किया गया है, जो Planck मास्क के आधार पर है, जिससे स्थानीय अग्रभूमि जटिलता के अनुकूल होने के लिए प्रत्येक क्षेत्र के लिए पैच गणनाओं के स्वतंत्र अनुकूलन की अनुमति मिलती है।

B. अनुकूलन और ग्रिड सर्च (Optimization & Grid Search)

• ग्रिड सर्च: विधि प्रत्येक स्पेक्ट्रल पैरामीटर के लिए पैच की संख्या ($K_{\beta_d}, K_{T_d}, K_{\beta_s}$) के लिए एक व्यापक ग्रिड सर्च करती है।
• चयन मीट्रिक: केवल मानचित्र प्रसरण (map variance) को न्यूनतम करने के बजाय (जो अंडर-फिटिंग को बढ़ावा देता है), इष्टतम विन्यास को $r$ के 68% ऊपरी विश्वास सीमा ($\hat{r} + \sigma(r)$) को न्यूनतम करके चुना जाता है। यह मीट्रिक स्पष्ट रूप से व्यवस्थित बायस (अग्रभूमि अवशेष) और सांख्यिकीय अनिश्चितता के बीच संतुलन बनाता है।
• CMB संरक्षण: पाइपलाइन यह सुनिश्चित करती है कि CMB संकेत बीजगणितीय रूप से संरक्षित रहे (यूनिटी रिस्पॉन्स), जिससे अनुकूलन को $r=0$ वाली सिमुलेशन पर बिना किसी हानि के किया जा सके।

C. कंप्यूटेशनल इंफ्रास्ट्रक्चर (FURAX और AdaTopK)

• JAX और GPU त्वरण: पाइपलाइन ऑटोमैटिक डिफरेंशिएशन और वेक्टराइजेशन (vmap) के लिए JAX का लाभ उठाती है, जो एकल GPU पर हजारों शोर वास्तविकताओं (noise realizations) और पैच विन्यासों के समवर्ती मूल्यांकन को सक्षम बनाती है।
• AdaTopK ऑप्टिमाइज़र: लेखकों ने एक कस्टम, JAX-नेटिव ऑप्टिमाइज़र विकसित किया है जिसे AdaTopK (Adaptive AdaBelief with Top-K Active Set) कहा जाता है।
  • यह मानक scipy TNC ऑप्टिमाइज़र को प्रतिस्थापित करता है।
  • यह लीनियर सॉल्वर के इम्प्लिसिट डिफरेंशिएशन (implicit differentiation) से बचने के लिए एक कस्टम बैकवर्ड पास (jax.custom_vjp) का उपयोग करता है, जिससे संख्यात्मक स्थिरता सुनिश्चित होती है।
  • यह चमकीले गैलेक्टिक प्लेन और मंद उच्च-अक्षांश आकाश के बीच अत्यधिक डायनेमिक रेंज को संभालने के लिए डायनेमिक स्टेट रीस्केलिंग का उपयोग करता है।
• प्रदर्शन: फ्रेमवर्क वर्कलोड को कई GPU में विभाजित करने के लिए वितरित कंप्यूटिंग (jax-grid-search) का उपयोग करता है।

3. मुख्य योगदान

1. GPU-त्वरित फ्रेमवर्क: FURAX का परिचय, जो एक मॉड्यूलर, JAX-नेटिव फ्रेमवर्क है जो स्पेक्ट्रल लाइकलीहुड मैक्सिमाइजेशन के लिए scipy TNC ऑप्टिमाइज़र्स की तुलना में $\sim$100$\times$ स्पीड-अप प्राप्त करता है।
2. AdaTopK ऑप्टिमाइज़र: एक नवीन अनुकूलन एल्गोरिथ्म जो विशेष रूप से उच्च-आयामी, बाउंडेड और शोर वाले समस्याओं के लिए डिज़ाइन किया गया है, जो लो सिग्नल-टू-नॉइज़ शासन में मानक क्वाजी-न्यूटन विधियों से बेहतर प्रदर्शन करता है।
3. डेटा-संचालित क्लस्टरिंग: निश्चित मल्टी-रेज़ोल्यूशन HEALPix ग्रिड के स्थान पर गोलीय K-मीन्स क्लस्टरिंग का उपयोग, जो स्पेक्ट्रल मापदंडों के लिए लचीले, डेटा-संचालित स्थानिक विभाजन की अनुमति देता है।
4. व्यापक मॉडल चयन: एक विशाल विन्यास स्थान ($>13,000$ पैच विन्यास 40 शोर वास्तविकताओं में) को खोजने की क्षमता, जो पहले गणनात्मक रूप से असंभव था।

4. परिणाम

फ्रेमवर्क को स्थानिक रूप से परिवर्तनशील अग्रभूमि (PySM3 d1s1 मॉडल) के साथ LiteBIRD-जैसे सिमुलेशन का उपयोग करके मान्य किया गया।

• गति: एक एकल NVIDIA H100 GPU पर AdaTopK, 40-कोर CPU नोड पर scipy TNC की तुलना में $\sim$100$\times$ तेज़ी से अभिसरण (converge) करता है (10,000 पैच के लिए)।
• इष्टतम विन्यास: ग्रिड सर्च ने पहचान की कि:
  • डस्ट स्पेक्ट्रल इंडेक्स ($\beta_d$) को लगभग पिक्सेल-स्तर के रेजोल्यूशन ($\sim$7,000–9,600 पैच) की आवश्यकता होती है।
  • डस्ट तापमान ($T_d$) को मध्यम रेजोल्यूशन ($\sim$300–2,500 पैच) की आवश्यकता होती है।
  • सिंक्रोट्रॉन इंडेक्स ($\beta_s$) अधिक स्मूथ है और इसके लिए कम पैच ($\sim$150–300) की आवश्यकता होती है।
• कॉस्मोलॉजिकल बाधाएं:
  • पिछले LiteBIRD अध्ययनों में उपयोग किए गए एक निश्चित, मैन्युअल रूप से ट्यून किए गए मल्टी-रेज़ोल्यूशन विन्यास की तुलना में, अनुकूलित K-means विन्यास ने $r$ की 68% ऊपरी सीमा को $\approx$30% कम कर दिया (1.1 $\times 10^{-3}$ से 7.4 $\times 10^{-4}$ तक)।
  • यह सुधार मुख्य रूप से व्यवस्थित बायस (अग्रभूमि अवशेषों) को कम करने से प्रेरित है, जबकि सांख्यिकीय अनिश्चितताएं तुलनीय बनी रहीं।
  • $r=0$ इनपुट के लिए, विधि ने $\hat{r} = 8.5 \times 10^{-5}$ को रिकवर किया।
  • $r=3 \times 10^{-3}$ इनपुट के लिए, विधि ने $\hat{r} = 3.1 \times 10^{-3}$ को रिकवर किया, जो पुष्टि करता है कि विधि CMB सिग्नल को कृत्रिम रूप से दबाती नहीं है।
• अवशेष (Residuals): अनुकूलित विन्यास ने निश्चित मल्टी-रेज़ोल्यूशन दृष्टिकोण की तुलना में कम मल्टीपोल ($\ell \lesssim 20$) पर व्यवस्थित अवशेषों को लगभग एक क्रम (one order) कम कर दिया।

5. महत्व

यह कार्य CMB घटक पृथक्करण में एक प्रतिमान परिवर्तन (paradigm shift) का प्रतिनिधित्व करता है, जो ह्यूरिस्टिक, निश्चित विभाजन से स्वचालित, डेटा-संचालित मॉडल चयन की ओर बढ़ता है।

• व्यवहार्यता: यह सिद्ध करता है कि बायस-वेरिएस ट्रेड-ऑफ स्पेस का व्यापक अन्वेषण अब संभव है, जिससे पैच विन्यासों के लिए मैन्युअल ट्यूनिंग की आवश्यकता समाप्त हो जाती है।
• भावी मिशन: यह विधि अगले पीढ़ी के CMB उपग्रहों (जैसे LiteBIRD) को अपने लक्ष्य संवेदनशीलता ($r < 0.001$) प्राप्त करने के लिए एक मजबूत मार्ग प्रदान करती है, जिससे अग्रभूमि बायस को सख्ती से कम किया जा सकता है।
• स्केलेबिलिटी: FURX फ्रेमवर्क को उच्च रिज़ॉल्यूशन ($N_{side} = 128-512$) और अधिक जटिल शोर/बीम मॉडल के लिए स्केल करने के लिए डिज़ाइन किया गया है, जो इसे भविष्य के ग्राउंड-बेस्ड और स्पेस-बेस्ड वेधशालाओं के लिए एक महत्वपूर्ण उपकरण बनाता है।
• ओपन साइंस: कोड और सिमुलेशन स्क्रिप्ट सार्वजनिक रूप से उपलब्ध हैं, जो समुदाय में पुनरुत्पादकता (reproducibility) और आगे के विकास को बढ़ावा देती हैं।

संक्षेप में, यह शोध स्थापित करता है कि GPU-त्वरित, डिफरेंशिएबल कंप्यूटिंग को अनुकूली क्लस्टरिंग के साथ जोड़ने से अग्रभूमि संदूषकों (foreground contaminants) को व्यवस्थित रूप से हटाकर मौलिक B-मोड डिटेक्शन की सटीकता में महत्वपूर्ण सुधार किया जा सकता है।