Three-dimensional Conditional Diffusion Models for Cosmological 21 cm Lightcone Emulation

Technical Summary: Three-Dimensional Conditional Diffusion Models for Cosmological 21 cm Lightcone Emulation

Problem Statement
The redshifted 21 cm signal from neutral hydrogen is a critical probe of the Cosmic Dawn and Epoch of Reionization (EoR). Interpreting future imaging data from facilities like HERA and SKA requires fast, physically reliable forward models across broad astrophysical parameter spaces. While semi-numerical simulations like 21cmFAST are standard, they are computationally expensive for large-scale parameter sweeps or simulation-based inference. Machine learning emulators, particularly generative models, offer a solution. However, extending these models from 2D image slices to full 3D lightcones introduces significant challenges. The addition of a third dimension (line-of-sight depth) drastically increases memory costs, forcing very small micro-batches that amplify optimization noise. Furthermore, the underlying voxel distribution of 21 cm brightness temperatures is highly skewed and long-tailed, making transform-space distortions and dynamic-range control critical for stable training and physical fidelity.

Methodology
The authors investigate conditional diffusion modeling for 3D 21 cm lightcone emulation using a dataset of 25,600 training lightcones generated by 21cmFAST.

• Data & Conditioning: The target fields are $64 \times 64 \times 1024$ voxel cubes spanning a redshift range of $z \in [7.51, 20.95]$. The models are conditioned on two astrophysical parameters: the minimum virial temperature of halos ($\log_{10} T_{\text{vir}}$) and the ionizing efficiency ($\zeta$).
• Architecture: A context-conditioned U-Net is employed, utilizing anisotropic downsampling/upsampling strides of $(2, 2, 4)$ to accommodate the elongated lightcone geometry. The model uses 128 base channels with channel multipliers $(1, 1, 2, 4)$ and incorporates timestep and condition embeddings.
• Preprocessing Pipeline: A central focus of the study is the preprocessing pipeline, defined as: $x_{\text{raw}} \to \text{primary transform} \to \text{optional squish} \to \text{diffusion} \to \text{inverse squish} \to \text{inverse transform}$. The authors compare four primary normalization transforms (Yeo–Johnson, Min–Max, Z-score, Arcsinh) and an optional dynamic-range compression step ("squish").
• Validation Design: The study uses 800 independent realizations per parameter point for validation across five distinct reference parameter sets. This design tests both the central tendency and the conditional diversity of the target distribution.
• Evaluation Metrics: Generated samples are evaluated using a hierarchy of diagnostics:
  1. Brightness-temperature ($T_b$) image slices (qualitative).
  2. Global signal $\langle T_b \rangle$ vs. redshift (mean evolution).
  3. Power spectrum (two-point spatial structure).
  4. Reduced scattering coefficients (higher-order non-Gaussian morphology).
     Performance is quantified using standard-deviation-normalized mean absolute error ($\text{MAE}_{\text{std}}$) as the primary ranking metric.

Key Contributions & Results

• Preprocessing Dominance: The study identifies preprocessing and dynamic-range control as the dominant factors governing training stability and physical fidelity in 3D settings, outweighing architectural depth or training duration once optimization is stabilized.
• Optimal Configuration: Among the tested configurations, Yeo–Johnson preprocessing combined with moderate amplitude compression ($A=0.1$) yields the most consistently favorable trade-off. This configuration outperforms Min–Max, Z-score, and Arcsinh transforms, as well as configurations without amplitude compression.
• Hierarchy of Fidelity: The emulation success follows a clear hierarchy:
  • Global Signal: The model achieves the strongest agreement, accurately capturing the timing and depth of absorption troughs.
  • Power Spectrum: Agreement is weaker, particularly at high-$k$ modes (small scales), where the model struggles to reproduce sharp boundaries.
  • Scattering Coefficients: Higher-order non-Gaussian morphology remains the most difficult to match. While the model captures broad structural patterns, measurable biases persist in multiscale morphology.
• Failure Modes: The primary limitation is the model's inability to perfectly reproduce the sharply concentrated, near-zero voxel population found in 21cmFAST. The diffusion model tends to broaden this peak, resulting in "smoother" ionized regions and small nonzero pixels in visually dark areas. This smoothing manifests as residual mismatches in the power spectrum at small scales and in higher-order statistics.
• Hyperparameter Sensitivity: Increasing model depth (residual blocks) did not yield monotonic improvements; in fact, the shallowest model (1 residual block) outperformed deeper variants, likely due to the noisy optimization dynamics induced by small micro-batches in 3D.

Significance and Claims
The authors position this work as a simulation-level baseline for three-dimensional 21 cm emulation rather than a complete end-to-end inference pipeline ready for survey data.

• Methodological Insight: The paper argues that in high-dimensional regimes, stable generation depends primarily on preprocessing and dynamic-range control. Visual plausibility alone is insufficient; morphology-sensitive diagnostics (like scattering coefficients) are necessary to detect biases that lower-order summaries might miss.
• Scope: The study is an interpolation test within the training parameter space. It does not yet address survey-facing challenges such as instrument response, foreground contamination, or uncertainty calibration for parameter estimation.
• Future Utility: The identified optimal configuration (Yeo–Johnson + $A=0.1$) provides a robust starting point for future studies aiming to incorporate realistic observational effects or perform simulation-based inference, provided that the limitations regarding small-scale morphology are acknowledged.