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

A novel, JAX-powered implementation of a parametric component-separation method for CMB polarization data is presented, explicitly designed to handle spatially varying foreground Spectral Energy Distributions (SEDs) by scanning thousands of pixel subset configurations to find the optimal trade-off between model complexity and residual systematic contamination.

The gist: This work introduces a JAX framework that accelerates the exploration of spatial clustering configurations in parametric component separation, achieving up to 100× speed-up over standard CPU optimizers while reducing the 68% upper limit on the tensor-to-scalar ratio by approximately 30% relative to fixed configurations.

Parametric Component Separation Model

The algebraic formulation follows the standard parametric formalism of Stompor et al. (2009), modeling observed data in each pixel as a linear combination of astrophysical components with additive Gaussian noise:

d p = A p (β p) s p + n p, where A(β) encodes spectral dependencies and s is the sky component vector. The negative log-likelihood under this model is given by Eq. (3). To estimate the sky components s, the generalized least squares solution yields the spectral likelihood of Stompor et al. (2009), which depends only on the spectral parameters β and is used to optimize their values before recovering sky amplitudes with Eq. (5).

Spatial Modeling via Patching

To account for spatial variability in foreground SEDs, the model generalizes by allowing β to vary across independent sky pixel subsets defined as spatially connected regions or patches. The method uses a spherical K-means clustering algorithm implemented via jax-healpy to partition the sky into disjoint patches where β p = β k for all p ∈ Pk. Crucially, this partitions the sky into roughly equal-area patches, and predefined templates such as Planck Galactic masks are used to divide the sky into regions with different levels of Galactic foreground contamination.

Grid Search over Patch Configurations

The approach treats the number of patches assigned to each spectral parameter (K βd, K Td, K βs) as discrete hyperparameters held fixed while optimizing the spectral parameters and recovering component amplitudes. A structured grid search is performed over G = ⋂G = ⋂G = G. For each configuration CK ∈ C, the spectral parameters are optimized by maximizing the spectral likelihood Lspec (Eq. 6), which depends on both parameter values in each patch and the spatial patch structure: ∀ CK ∈ C: β k∗ = arg max β k Lspec (β k, CK).

Selection Metrics for Patch Configurations

The optimal configuration is selected by evaluating metrics that quantify residual contamination and cosmological performance. The primary selection criterion is the upper limit of the 68% confidence interval on the tensor-to-scalar ratio, formally rˆ + σ(r), which explicitly sums systematic bias (driven by foreground residuals) and statistical uncertainty. This metric minimizes both excessive bias and excessive variance, as minimizing rˆ alone favors maximal patch counts that inflate statistical noise.

Computational Infrastructure and Optimization

The pipeline is implemented within FURAX (Chanial et al. 2026), a JAX-powered framework designed for end-to-end differentiability and GPU acceleration. It leverages intra-device parallelism via JAX’s vmap transformation to vectorize operations over noise realizations. The central task of maximizing Lspec is solved using the custom AdaTopK optimizer, which is a JAX-native active-set algorithm that outperforms standard scipy TNC on CPU by up to ∼100×, demonstrating improved robustness in low-SNR regions.

Sky Region Partitioning and Results

The analysis partitions the sky into three disjoint Galactic regions: high-latitude (hi-lat), mid-latitude (mid-lat), and low-latitude (low-lat) based on Planck Galactic masks. The optimal patch counts vary by region, revealing that the dust spectral index requires near-pixel resolution, while the synchrotron spectral index is well described by fewer patches. The optimized configuration yields an estimate of rˆ = 8.5 × 10−5 with a 68% confidence interval of [−2.7 × 10−4, +6.5 × 10−4], representing a ∼30% reduction in the upper limit on r compared to fixed, manually tuned configurations.

Future Developments

Future work aims to move towards more flexible pixel subsets by exploring discontinuous pixel subsets through grouping multiple K-means clusters or employing more flexible clustering algorithms like Support Vector Machine clustering. This approach seeks to reduce statistical residuals without increasing systematic residuals, although the paper notes that a naive binning approach does not perform well compared to the optimized K-means configuration. Additionally, extending the framework to incorporate instrumental systematics and higher resolutions is planned.

Data Availability

The code used for generating results is publicly available on GitHub at https://github.

As a fastidious researcher, I have analyzed this paper on GPU-accelerated JAX framework for robust parametric component separation of CMB polarization data. The core innovation lies in efficiently exploring the space of spatially varying foreground spectral Energy Distributions (SEDs) using K-means clustering and an optimized JAX optimizer (AdaTopK).

Based on the methodology, here are specific improvements that can be implemented in AI systems derived from this research, along with the capabilities these improved systems would possess:

)

1. Improvement: Develop a highly efficient, GPU-accelerated differentiable pipeline for modeling and separating spatially varying physical components in complex datasets (like astrophysical signals).

2. Capability: The resulting system can perform high-dimensional parameter inference across massive, heterogeneous observational data (e.g., multi-frequency satellite observations) while robustly accounting for spatial variations in the underlying physical processes (foreground SEDs).

3. Improvement: Integrate a novel optimization algorithm, AdaTopK, which is JAX-native and designed to handle high-dimensional, bounded, and noisy objective functions with superior convergence speed and robustness compared to traditional quasi-Newton methods (like SciPy TNC).

4. Capability: This allows the AI system to rapidly converge on optimal physical models (e.g., finding the best combination of dust temperature, spectral index, and synchrotron parameters) even in low signal-to-noise regions or when dealing with complex, non-convex likelihood landscapes common in real-world data analysis.

5. Improvement: Implement a flexible spatial partitioning strategy using spherical K-means clustering (instead of fixed HEALPix superpixels), allowing the system to automatically define pixel subsets based on shared spectral parameter values, tailored to local foreground complexity (e.g., dividing the sky into high-latitude, mid-latitude, and low-latitude regions).

6. Capability: The AI system gains the ability to perform spatially adaptive regularization. It can dynamically adjust its modeling resolution—using fine patches in high-foreground regions and coarse patches in low-noise regions—to achieve a superior bias-variance trade-off, leading to significantly tighter constraints on cosmological parameters like the tensor-to-scalar ratio (r).

7. Improvement: Replace simple map variance metrics with a sophisticated, integrated selection criterion that minimizes the 68% upper limit on the tensor-to-scalar ratio, defined as the sum of systematic bias and statistical uncertainty.

8. Capability: The system moves beyond simply minimizing noise in the resulting maps; it is optimized directly for cosmological performance. This ensures that the final derived parameter estimates (like r) are not just statistically small but also free from significant, model-induced systematic errors, leading to more reliable scientific conclusions about primordial gravitational waves.

9. Improvement: Implement a modular, end-to-end framework (FURAX) built on JAX for composable linear operators and differentiable algebraic operations, allowing for easy integration of new physical models (e.g., beam convolution or correlated noise).

10. Capability: The system becomes highly adaptable to next-generation instrumentation and more realistic data complexities. It can seamlessly incorporate instrumental effects (like beam asymmetries or gain drift) into the forward model, enabling rigorous testing against real satellite data and providing a framework for future mission planning.

11. Improvement: Develop a methodology for exploring discontinuous pixel subsets by binning recovered spectral parameter templates, which reduces statistical residuals without drastically increasing systematic ones.

12. Capability: The system can effectively manage the trade-off between statistical noise (from many patches) and systematic bias (from poor spatial grouping), leading to cleaner residual maps and more accurate recovery of the primordial signal, even when dealing with complex foreground structures that do not fit a simple uniform model.

We present a novel, JAX-powered implementation of a parametric component-separation method for CMB polarization data, explicitly designed to handle spatially varying foreground Spectral Energy Distributions (SEDs). The approach models this variation across the sky by grouping sets of pixels that share common foreground spectral parameters, scanning over thousands of such configurations to evaluate the trade-off between model complexity and residual systematic contamination. Built within the FURAX framework -- a JAX-powered environment for CMB data analysis -- our pipeline extends the fgbuster parametric formalism. It enables fully vectorized, GPU-accelerated evaluation of the spectral likelihood, map reconstruction, and diagnostic metrics across tens of thousands of pixel subset configurations, noise realizations, and sky regions. Our implementation achieves up to about 100 times speed-up over the scipy TNC optimizer used in fgbuster when running on GPUs, as well as giving more robust results. When applied to LiteBIRD-like simulations with spatially varying foreground SEDs, our optimized K-means configuration reduces the 68% upper limit on the tensor-to-scalar ratio r by about 30% relative to a fixed, previously derived multi-resolution configuration, while maintaining competitive statistical uncertainties.