Explaining Neural Networks on the Sky: Machine Learning Interpretability for Cosmic Microwave Background Maps

1. Problem Statement

The Cosmic Microwave Background (CMB) is the primary probe of the early Universe. While the standard $\Lambda$CDM model successfully describes temperature and polarization anisotropies via angular power spectra ($C_\ell$), this approach relies on second-order statistics. Consequently, it compresses data, discarding phase information and higher-order correlations.

• The Challenge: Many theoretical extensions to $\Lambda$CDM, specifically primordial feature models (e.g., oscillatory deviations in the primordial power spectrum caused by sharp turns in field space or multi-field interactions), leave subtle, non-Gaussian, or spatially localized imprints that are often obscured in 1D harmonic-space analyses.
• The Gap: Direct analysis of CMB maps using Machine Learning (ML) offers a path to recover this lost information. However, ML models are often "black boxes," making it difficult to verify if they are learning physical signals or spurious correlations (e.g., artifacts from galactic masking or noise). There is a critical need for interpretable ML frameworks that can operate directly on map-level data to distinguish between standard and feature models.

2. Methodology

The authors propose an end-to-end pipeline that combines simulation, hybrid neural network architecture, and post-hoc interpretability.

A. Data Simulation and Pre-processing

• Simulations: Full-sky CMB temperature ($T$) and polarization ($Q, U$) maps were generated using CAMB to compute theoretical power spectra for both $\Lambda$CDM and a feature model.
  • Feature Model: Defined by a modulation of the primordial power spectrum: $P_R(k) = P_{R,0}(k)[1 + A_{lin} \sin(\omega_{lin} \frac{k}{k_\star} + \phi)]$.
  • Instrumental Realism: Maps were convolved with a Planck-like beam (FWHM $\sim$ 5 arcmin), added with anisotropic noise, and processed with a Galactic mask to mimic real observational conditions (removing foregrounds).
  • Resolution: HEALPix $N_{side} = 256$ ($\sim 8 \times 10^5$ pixels), balancing angular resolution with computational cost.
• Pre-processing: Instead of standard Z-score normalization, the authors used per-map standardization ($X' = (X - \mu)/(\sigma + \epsilon)$) to ensure classification relies on morphological structure rather than global amplitude variations.

B. Hybrid Neural Network Architecture

To address the high dimensionality of masked maps and the difficulty of optimizing raw pixel data, a Hybrid PCA-MLP architecture was employed:

1. PCA Layer (Feature Extraction): A Principal Component Analysis (PCA) transformation is implemented as a fixed, non-trainable first layer. It projects the high-dimensional pixel vectors onto the top $k=200$ principal components.
   • Rationale: This reduces dimensionality, removes redundancy, and concentrates the dominant variance (including the feature signal) while filtering out noise. It improves the conditioning of the optimization problem, which failed when training CNNs directly on raw pixels.
2. MLP Branch (Classification): The reduced feature vector is fed into a Multi-Layer Perceptron (MLP) with ReLU activations, hidden layers (64–128 neurons), and dropout.
   • Separate pipelines were trained for Temperature ($T$) and Polarization ($U, Q$) maps.
   • Loss Function: Binary Cross-Entropy optimized via the Adam algorithm.
   • Regularization: Early stopping based on validation AUC (Area Under the Curve).

C. Interpretability (SHAP)

To validate that the network learns physical signals, the authors applied SHapley Additive exPlanations (SHAP):

• Algorithm: Partition SHAP was used to handle spatial correlations in HEALPix maps efficiently.
• Goal: To generate pixel-wise attribution maps ($\phi_i$) showing which regions of the sky contribute most to the classification decision. This verifies if the model focuses on cosmological residuals or masking artifacts.

3. Key Contributions

1. Map-Level Analysis Framework: Demonstrated a successful pipeline for classifying cosmological models directly on $T, Q, U$ maps, bypassing the information loss inherent in power spectrum compression.
2. Hybrid PCA-NN Architecture: Introduced a novel architecture where PCA acts as a fixed, non-trainable input layer. This solved the optimization instability of training deep networks on raw, high-dimensional, masked CMB data.
3. Interpretability Validation: Applied SHAP to CMB maps to provide a "spatially resolved" verification of the model's logic, proving the network detects global oscillatory signatures rather than local noise or mask edges.
4. Open Science: Released the full pipeline (simulation, training, and interpretability tools) via the SkyExplain GitHub project and Zenodo.

4. Results

• Classification Performance:
  • The PCA-MLP pipeline achieved near-perfect classification accuracy (AUC $\approx$ 1.0) in distinguishing $\Lambda$CDM from feature models with amplitudes $A_{lin} \in [0.01, 0.06]$ and frequencies $\omega_{lin} \in [5, 25]$.
  • Ablation Study: Removing the PCA step caused performance to drop to near-random levels, confirming that the discriminative signal is inaccessible in raw pixel space without dimensionality reduction.
• Interpretability Findings:
  • Global Distribution: SHAP attribution maps showed that the network's decision-making relies on globally distributed spatial signatures across the high-latitude sky, consistent with the theoretical expectation of primordial oscillations.
  • Robustness to Masking: The attribution patterns remained consistent under Galactic masking, with no significant "edge effects" near the mask boundaries.
  • Null Test: For pure $\Lambda$CDM maps, the SHAP values were near-zero (featureless), confirming the model does not overfit to the Gaussian random field or noise.
  • Cross-Correlation: The model effectively utilized correlations between $T$, $Q$, and $U$ components.

5. Significance

This work represents a paradigm shift in CMB analysis by moving from summary statistics to interpretable, map-level deep learning.

• Physical Trust: By using SHAP, the authors bridge the gap between "black box" AI and physical cosmology, demonstrating that the model's decisions are driven by the specific cosmological signatures of interest.
• Future Probes: The framework is scalable and adaptable. It offers a powerful tool for next-generation CMB experiments (e.g., CMB-S4, Simons Observatory) to detect subtle, non-Gaussian, or localized inflationary signals that traditional likelihood-based methods might miss due to masking and deconvolution artifacts.
• Methodological Standard: It establishes a new standard for transparency in ML-based cosmological discovery, ensuring that detected anomalies are physically grounded rather than statistical artifacts.