ResUNet-CMB: Deep Learning for CMB Anisotropies

• ResUNet-CMB is a deep convolutional neural network that jointly reconstructs secondary CMB anisotropies such as gravitational lensing, patchy reionization, and polarization rotation.
• The network leverages a U-Net-style encoder-decoder with both residual and skip connections to enable efficient, bias-reduced, end-to-end map-level inference.
• Extensive simulations show that ResUNet-CMB achieves lower noise and better signal disentanglement than quadratic estimators, making it ideal for next-generation CMB surveys.

ResUNet-CMB is a deep convolutional neural network architecture developed for the simultaneous, map-level reconstruction of multiple secondary anisotropies in cosmic microwave background (CMB) polarization data, most notably gravitational lensing, patchy reionization, and anisotropic cosmic polarization rotation. Its design leverages a U-Net-style encoder–decoder backbone with both residual and U-shaped skip connections, enabling efficient end-to-end inference of fields such as lensing convergence (κ), patchy optical-depth modulation (τ), anisotropic rotation (α), and the primordial unlensed E-mode (E). ResUNet-CMB is positioned as a direct competitor to—and, in several respects, a surpassing alternative to—standard quadratic estimators, with greatly reduced bias and improved noise properties at the sensitivity of upcoming CMB surveys (Guzman et al., 2021, Guzman et al., 2021).

1. Network Architecture and Computational Graph

The ResUNet-CMB architecture is a U-shaped fully convolutional autoencoder, augmented by residual connections both locally (within the conv-block streams) and globally (across encoder–decoder levels). The standard implementation operates on flat-sky Stokes Q and U maps, typically of size 128×128 or larger, and predicts one or more output maps: κ(n), τ(n), α(n), and E(n), with output branch configuration matched to the reconstruction task.

• Convolutional Blocks: Each block comprises dropout (rate ~0.3), 2D convolution (typically 3×3 or 5×5, “same” padding), SELU or ReLU nonlinearity, and batch normalization.
• Encoder: Progressive spatial downsampling is achieved via stride-2 convolutions, doubling channel dimension at each depth.
• Decoder: Upsampling mirrors the encoder, employing nearest-neighbor resize or transposed convolution, followed by convolutional processing.
• Residual (Intra-stream) Connections: Every two conv-blocks are connected by an identity addition path, with linear convolution for shape/width alignment if needed, ensuring stable gradient flow and regularizing deeper layers (Guzman et al., 2021, Guzman et al., 2021, Heinrich et al., 2022).
• Skip (U-Net) Connections: At regular intervals (e.g., every three (ResUNet-CMB) or matching encoder/decoder depths), encoder feature maps are concatenated with decoder activations at corresponding spatial resolution, restoring high-frequency (fine spatial) information.
• Output Heads: For multi-field tasks (e.g., joint κ, τ, E, α reconstruction), final layers are 1×1 convolutions with linear activations, producing continuous-valued maps for each parameter.
• Receptive Field: For typical configurations, the receptive field is >100×100 pixels (~4°).
• Parameter Count: ~5–6 million for standard setups.

2. Mathematical and Physical Formulation

ResUNet-CMB is designed to invert the complex, non-linear mapping from primordial CMB observables (T, Q, U) to the observed maps that have experienced multiple anisotropic distortions:

• Patchy Reionization:

$$X_\mathrm{mod}(n) = X_\mathrm{prim}(n) \, e^{-\tau(n)}$$

where τ(n) represents anisotropic optical depth modulation from inhomogeneous reionization (Guzman et al., 2021).

• Gravitational Lensing:

$$X_\mathrm{obs}(n) = X_\mathrm{mod}(n + \nabla\phi(n))$$

with κ(n) = ½∇²φ(n) as the convergence.

• Cosmic Polarization Rotation:

$$[Q' \pm iU'](n) = e^{\pm 2i\,\alpha(n)}\,[Q \pm iU](n)$$

where α(n) is the local polarization rotation angle (Guzman et al., 2021).

• Training Objective:

The canonical loss is the mean-squared error, summed over all output fields:

$$L = \left\langle \|\hat{\kappa} - \kappa\|^2 + \|\hat{\tau} - \tau\|^2 + \|\hat{E} - E\|^2 + \|\hat{\alpha} - \alpha\|^2 \right\rangle$$

with averaging over pixels and training batch.

3. Training Pipeline and Data Simulation

• Simulation Inputs: Training utilizes CMB realizations generated by CAMB (ΛCDM cosmology), anisotropic τ field (e.g., from the Roy et al. model for patchy reionization), lensing potentials φ(n) as Gaussian random fields, and (for α) rotation fields with specified scale-invariant spectra.
• Forward Model: Sequential application of modulation (reionization), rotation (if α present), lensing remapping, Gaussian beam convolution, and additive isotropic white noise at target noise levels (Δ_T, Δ_P).
• Data Volume: Typical runs use 70,000 training realizations, up to ~7,000 test maps per configuration, and multiple noise settings (Δ_T = {0, 0.2, 1, 2} μK-arcmin).
• Normalization: Inputs and outputs are standardized using the mean and variance of each map type in the training set.
• Null Maps: A significant fraction (~20%) are simulations with κ=τ=α=0 to prevent hallucination of spurious signals—a key ingredient for unbiased null tests (Guzman et al., 2021, Guzman et al., 2021).
• Optimization: Adam optimizer (default β), initial learning rates ~0.25, decay on plateau, batch size = 32. Early stopping based on validation loss stability.

4. Quantitative Performance and Comparison with Other Estimators

• Reionization Field Recovery: For τ(n), the ResUNet-CMB achieves N_ℓ^{ττ} (noise power) factors of 2–3 lower than the quadratic estimator (QE) around multipoles ℓ~500, and matches or surpasses bias-hardened/iterative QE at low noise. Lensing-induced bias, which strongly contaminates QE-based τ reconstructions, is suppressed to negligible levels by ResUNet-CMB across ℓ≲1000 (Guzman et al., 2021).
• Lensing Reconstruction: κ(n) reconstruction noise N_L^{κκ} with ResUNet is 50–70% below QE for 200<L<2000, comparable to iterative maximum-likelihood estimator ("MLE EB N_0") (Caldeira et al., 2018, Li et al., 2022).
• Rotation Field α(n): When trained for rotation, the ResUNet-CMB N_ℓ^{αα} noise is reduced by factors of 3–5 compared to standard QE on large scales, closely tracking the optimal iterative estimator for low instrumental noise; performance for τ, κ is only modestly affected by including α in the output head (Guzman et al., 2021).
• Signal Disentanglement: Cross-correlation coefficients (e.g., r_{κ,τ̂} vs. r_{τ,τ̂}) demonstrate the model’s ability to avoid leaking lensing into τ or vice versa, surpassing QE which is lensing-dominated in mixed scenarios.
• Robustness: Null tests with pure noise inputs result in output auto-spectra many orders of magnitude below the true signals, indicating strong resistance to hallucination (Guzman et al., 2021).
• Speed: Map-level inference is ~0.02 s/patch, suitable for large-scale Monte Carlo or Bayesian cosmological applications.
• Limitations: Performance on small angular scales (ℓ≳2000) and at high noise levels (>1 μK-arcmin) degrades, with reconstruction noise and cross-correlation coefficients approaching zero for τ or α. The architecture may miss very fine spatial structure due to its receptive field and the smoothing tendency of the decoder.

5. Extension to Joint and Multi-field Reconstruction

ResUNet-CMB supports joint, multi-branch architectures, reconstructing multiple secondary anisotropies simultaneously—lensing (κ), patchy reionization (τ), polarization rotation (α), and E modes (E):

• Simultaneous Learning: The end-to-end mapping enables the network to learn non-linear disentanglement of physical effects, effectively internalizing the role of several families of optimal quadratic estimators and their iterative combinations (Guzman et al., 2021).
• Stability: Adding additional output branches for τ or α increases reconstruction noise for κ by <10% on relevant scales.
• Comparison to GAN/cGAN Approaches: While vanilla ResUNet-CMB provides high-fidelity, low-noise reconstructions on large scales, it can underestimate small-scale power and structure (particularly for cluster-scale convergence); cGAN approaches can halve high-ℓ power bias and recover sharper small-scale features at the cost of noise (Parker et al., 2022, Li et al., 2022).
• Adaptability: The architecture is readily adaptable to broader input spaces (e.g., inclusion of temperature maps, multi-frequency data, or foregrounds) via input head modification or additional branches.

6. Applicability to Next-Generation CMB Surveys and Forward-looking Directions

ResUNet-CMB is designed to be compatible with the data volumes, resolutions, and noise levels expected for Stage-3/4 ground-based CMB experiments (e.g., Simons Observatory, CMB-S4). It offers inference-level speed for map reconstructions, integrates robustly into cosmological likelihood analyses, and is extensible to new science cases:

• Survey Optimizations: Near-optimality is achieved at Δ_T ≲ 1 μK-arcmin, typical for the next generation of probes.
• Pipeline Integration: Plug-and-play inference speed and ease of wrap-around make ResUNet-CMB suitable for use in full-survey Monte Carlo, forward simulation, or Bayesian pipelines.
• Future Architectures: Extensions to all-sky (spherical) U-Net designs, explicit uncertainty quantification (Bayesian neural nets, variational layers), hybrid physical–neural models, and test-time adaptation to foregrounds or systematics are straightforward given the modularity of the approach (Sudevan et al., 27 Jun 2024).
• Null Testing and Calibration: Null maps and simulation-driven calibration are essential for credible inference, as demonstrated in multiple works (Guzman et al., 2021, Guzman et al., 2021).
• Current Limitations and Open Challenges: Small scale and high-noise regime performance, rigidity to spectral priors, and lack of analytic uncertainty estimates in basic deterministic models remain active areas for improvement (Li et al., 2022, Parker et al., 2022).

7. Impact and Comparative Summary

The introduction of ResUNet-CMB demonstrates that deep convolutional networks, specifically those with encoder–decoder and residual/U-Net structure, are capable of saturating or even surpassing the information bounds of traditional quadratic estimators for joint reconstruction of secondary CMB anisotropies in the low-noise regime relevant for next-generation cosmology. The architecture enables direct, end-to-end map-level disentanglement of lensing, reionization, and rotation signals, with negligible bias under null conditions, strong control of signal leakage, robust performance to noise, and compatibility with modern simulation-based inference pipelines (Guzman et al., 2021, Guzman et al., 2021, Guzman et al., 22 Dec 2025). This positions ResUNet-CMB as a foundation for continued advances toward precision CMB cosmology in an era where secondary anisotropies and their non-Gaussian structure play a dominant role in constraining fundamental physics parameters.