Robustness of Neural Networks for CMB Polarization Foreground Removal

The gist The generalization hypothesis suggests that training on more statistically complex foreground models leads to improved generalization across different scenarios for CNN-based methods in CMB polarization foreground removal.

Introduction and Motivation

The detection of Cosmic Microwave Background primordial B-mode polarization would constitute a “smoking gun" signal of primordial gravitational waves, but this measurement requires accurate removal of polarized Galactic foregrounds to avoid systematic biases when estimating the tensor-to-scalar ratio The paper focuses on testing the applicability of CNNs for foreground cleaning by systematically studying their generalization properties across various Galactic Foreground Models (FMs)

Simulation Setup

The simulations involve three components: lensed CMB polarization maps, instrumental noise realization maps, and polarized Galactic foregrounds. The analysis uses a partial-sky approach centered at (RA, DEC) = (316◦, −56◦), corresponding to a sky fraction of fsky ≈ 0.015. For the CNN-based method, simulated training, validation, and testing sets were generated with 1000, 300 and 300 polarization maps respectively. The ground truth is defined by computing the pseudo-Cls of the lensed CMB polarization maps using Namaster to estimate the pseudo-Cl, including B-mode purification.

Foreground Models and Complexity

The study considers three foreground models for training CNNs: "GP based on Gaussian fluctuations of SED parameters, RP based on rotating the sky before extracting the sky patch, and PySM3 d11s6". Statistical complexity is quantified using variance, skewness, and Shannon entropy to define a hierarchy among FMs. The Rotating Patches (RP) model is identified as the most statistically complex FM based on these metrics.

CNN Method and Training

The CNN-based method assumes a foreground model during both training and testing. Data pre-processing involves transforming Healpy sky maps into 2D blocks of shape (Nsize, Nsize) for CNN input, with Nsize = 128 for simulations with Nside = 256. The training procedure involves selecting a network architecture (UNet-type, UB, or L310), using a loss function consisting of the Mean Absolute Error (MAE) and a physics-guided term involving the Fast Fourier Transform (FFT).

Generalization Hypothesis and Results

The generalization hypothesis posits that CNNs trained on the RP model are expected to exhibit an improved performance if tested on FMs different from those used during training. Performance is quantified using the dimensionless ratio ratX,rl[M1, M2] = CRec,X,rl[M1, M2] / CTrue,X,rl[M2], and summary statistics like the median bias biasXl[M1, M2] which measures accuracy. Results show that L3-RP performs better on GP than L3-d11s6 when tested on d11s6, and L3-GP outperforms L3-d11s6 when tested on RP. The largest systematic deviations occur for L3-GP when tested on d11s6, representing the only case where the E-mode bias exceeds 1σ.

Conclusion

Training assuming the statistically most complex model improves foreground removal performance, and CNNs trained on RP achieve better performance when tested on FMs not used during training. This result supports the generalization hypothesis that training on more statistically complex FMs leads to improved generalization across different scenarios. The interpretation of generalization in terms of these metrics provides a new framework for studying it quantitatively.

Acknowledgements

This work was supported by the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http://iaifi.org/). L.G.B. would like to thank Guillermo E. Perna for helpful discussions on the Shannon entropy interpretation.

References

[1] A. H. Guth, Physical Review D 23, 347 (1981).

[2] M Tegmark and G. Efstathiou, Monthly Notices of the Royal Astronomical Society 281, 1297 (1996).

[3] P. A. Ade, Z. Ahmed, M. Amiri, D. Barkats, R. B. Thakur, C. Bischoff, D. Beck, J. Bock, H. Boenish, E. Bullock et al., Physical review letters 127 151301 (2021).

[4] M Tristram, A J Banday K M Gorski R Keskitalo C R Lawrence K J Andersen R B Barreiro J Borrill L P L Colombo H K Eriksen R Fernandez-Cobos T S Kisner E Martinez-Gonzalez B Partridge D Scott T Svalheim and I Wehus, Physical Review D 105 (2019).

[5] M Tegmark and G. Efstathiou, Monthly Notices of the Royal Astronomical Society 281, 1297 (1996).

[6] W Hu M Hedman and M Zaldarriaga, Physical Review D 67 043004 (2003).

[7] S Stellati F Piacentini S Micheli A Novelli F Columbro A Coppolecchia P de Bernardis S Masi M Najafi A Occhiuzzi et al., arXiv preprint arXiv:2509.03564 (2025).

[8] A Lewis and A Challinor, Physics Reports 429 1 (2006).

[9] S Belkner J Carron L Legrand C Umiltà C Pryke C Bischoff C-S Collaboration et al., The Astrophysical Journal 964 148 (2024).

[10] P A Ade N Aghanim Z Ahmed R W Aikin K D Alexander M Arnaud J Aumont C Baccigalupi A J Banday D Barkats R B Thakur C Bischoff J Bock H Boenish E Bullock et al., Physical review letters 114 101301 (2015).

[15] R Stompor S Leach F Stivoli and C Baccigalupi, Monthly Notices of the Royal Astronomical Society 392 216 (2009).

[67] C E Shannon, The Bell system technical journal 27 379 (1948).

--- Page 1 ---

Robustness of Neural Networks for CMB Polarization Foreground Removal Luca Gomez Bachar Cora Dvorkin and Alberto Daniel Supanitsky The detection of Cosmic Microwave Background primordial B-mode polarization would constitute a “smoking gun" signal of primordial gravitational waves. However, this measurement requires accurate removal of polarized Galactic foregrounds to avoid systematic biases when estimating the tensor-to-scalar ratio Methods based on Machine Learning techniques (ML), such as Convolutional Neural Networks (CNNs), have recently been proposed as alternative foreground cleaning techniques, but their applicability to real data relies on their ability to generalize beyond the models assumed during training In this work, we focus on a variety of foreground models (FMs) used for training and conduct a systematic study of the generalization properties of a CNN-based method We show that training on the more complex FMs reduces bias and improves precision when testing on unseen FMs, whereas training on the simplest model could introduce systematic errors These results evidence that a lack of generalization is a relevant source of systematic uncertainty, and emphasize the importance of understanding the impact of the models assumed during training in ML-based methods before applying them to real data.

--- Page 2 ---

In this work, we study the generalization properties of a CNN-based foreground-cleaning method applied to CMB polarization, adopting as a case study a CMB-S4-like small aperture telescope (SAT). Rather than assuming that neural networks automatically generalize across different FMs, we explicitly test this assumption by training CNNs on simulations generated from specific template-based Galactic FMs and evaluating their performance on other models not seen during the training. To this end, we introduce performance metrics to quantify both the accuracy and precision of the reconstructed CMB signal and to identify potential biases and uncertainty arising from the lack of generalization. To interpret these results, we quantify statistical properties (variance, skewness, and Shannon entropy) of the FMs considered in this work. We show that these statistics serve as a measure of foreground complexity. In this sense, generalization emerges not merely as a machine learning concern but as a physically relevant property for controlling the systematic uncertainty in the estimation of cosmological parameters. This paper is organized as follows. In Section II we describe the simulation pipeline used to generate CMB polarization maps, instrumental noise, and Galactic foregrounds. In Section III we present the foreground-cleaning methods considered in this work, including the traditional techniques and the CNN-based approach.

1. Bold generalization of CNN architectures to include batch normalization layers (UB architecture) is beneficial for controlling gradient norms, as when batch normalization is added (UB architecture), these large gradient values disappear. This suggests that incorporating specific regularization techniques can stabilize training and improve convergence for neural networks in this domain.

2. Explicitly using statistical complexity hierarchy to select the most relevant foreground model (RP) for training is crucial, as we show that training on the more complex FMs reduces bias and improves precision when testing on unseen FMs. This allows AI systems to be trained on the most challenging realistic scenarios rather than simpler ones, directly mitigating a lack of generalization is a relevant source of systematic uncertainty.

3. Development of quantitative performance metrics like the bias Generalization Factor (bias-GF) allows for direct comparison between networks trained on different FMs when tested on a third model. This enables researchers to determine which training strategy leads to better outcomes, as indicated by the finding that the CNNs trained with RP achieve better performance when tested on FMs not used during training.

4. Implementation of metrics like Nσ Xl[M1, M2] provides a stringent test of reliability by calculating the number of standard deviations that the median ratio deviates from unity for the L3 architecture when tested on FMs not used during training. This allows AI systems to be evaluated against a threshold where a value greater than one indicates that the 68% probability interval does not encompass the true CMB power spectrum.