HaloFlow II: Robust Galaxy Halo Mass Inference with Domain Adaptation

1. Problem Statement

Precise measurement of dark matter halo masses ($M_h$) is critical for cosmology and galaxy evolution studies. However, $M_h$ is not directly observable. Traditional methods (gravitational lensing, abundance matching, group dynamics) have limitations regarding resolution, applicability to low-mass halos, or the ability to provide calibrated posterior distributions for individual galaxies.

Simulation-Based Inference (SBI) offers a solution by using machine learning to map galaxy observables (images) to physical parameters. The authors' previous work, HaloFlow, used Neural Posterior Estimation (NPE) with normalizing flows to infer $M_h$ from synthetic galaxy images with high precision.

The Core Challenge: HaloFlow suffers from domain shift. When a model trained on one cosmological simulation (e.g., IllustrisTNG) is applied to a different simulation (e.g., EAGLE or SIMBA) or real observational data, the underlying galaxy formation physics differ. This "out-of-domain" application causes the model to produce biased and overconfident posteriors, rendering the inference unreliable. The paper aims to solve this by integrating Domain Adaptation (DA) to make the inference robust across different simulation domains.

2. Methodology

The authors introduce HaloFlowDA, an extension of HaloFlow that incorporates unsupervised domain adaptation techniques before the SBI step.

Data

• Simulations: Three state-of-the-art hydrodynamical simulations with distinct subgrid physics: IllustrisTNG (TNG50/TNG100), EAGLE, and SIMBA.
• Images: Forward-modeled synthetic images mimicking the Hyper Suprime-Cam Subaru Strategic Program (HSC-SSP). These include realistic dust attenuation, noise, and point spread functions.
• Features: A vector $X$ comprising photometric magnitudes, morphological parameters (Sérsic index, effective radius), and non-parametric statistics (CAS, Gini, $M_{20}$).
• Target: Central galaxies at $z=0.1$ with stellar mass $M_* > 10^{10} M_\odot$.

Domain Adaptation Strategies

The authors test two DA methods to map raw features $X$ into compressed, domain-invariant representations $\tilde{X}$:

1. Domain-Adversarial Neural Networks (DANN):

   • Uses an adversarial setup with three components: a feature extractor, a regressor (predicting $M_*, M_h$), and a domain classifier.
   • A Gradient Reversal Layer (GRL) is applied to the domain classifier. During backpropagation, the gradient is reversed, forcing the feature extractor to learn features that maximize regression accuracy while minimizing the ability of the classifier to distinguish which simulation the data came from.
   • The compressed representation $\tilde{X}$ is the 32-dimensional latent output of the feature extractor.

2. Maximum Mean Discrepancy (MMD):

   • A non-adversarial approach that directly minimizes the statistical distance between feature distributions of source and target domains.
   • It embeds latent features into a Reproducing Kernel Hilbert Space (RKHS) using a Gaussian Radial Basis Function (RBF) kernel.
   • The loss function minimizes the squared MMD distance between the latent features of the target domain and the source domains.
   • The compressed representation $\tilde{X}$ is the 2-dimensional regressor output ($M_*, M_h$).

HaloFlowDA Pipeline

1. Feature Compression: Train a DA network (DANN or MMD) on labeled source domains and unlabeled target domain data to produce $\tilde{X}$.
2. SBI Training: Train the HaloFlow NPE (normalizing flow) using the compressed features $\tilde{X}$ instead of raw features $X$.
3. Inference: Apply the trained NPE to the target domain to generate posterior distributions $p(M_*, M_h | \tilde{X})$.

3. Key Contributions

• HaloFlowDA Framework: A novel integration of unsupervised domain adaptation with simulation-based inference for galaxy halo mass estimation.
• Systematic Evaluation: A comprehensive benchmark across all pairwise combinations of TNG, EAGLE, and SIMBA simulations to quantify domain shift effects.
• Comparative Analysis of DA Methods: A detailed comparison between adversarial (DANN) and distance-based (MMD) adaptation, revealing that MMD is significantly more stable and robust in this context.
• Prior Correction Improvements: Refined the HaloFlow prior correction method by fitting a Schechter profile to the halo mass function and truncating posterior tails to ensure numerical stability and uniform priors.

4. Results

• Impact of Domain Shift: Without DA, applying a model trained on TNG to EAGLE resulted in a normalized residual metric ($\beta$) of 1.66 (indicating severe bias and overconfidence), compared to 0.58 for in-domain testing. Coverage dropped from ~77% (in-domain) to ~48% (out-of-domain).
• Performance of HaloFlowDA:
  • Both DANN and MMD significantly reduced bias and improved calibration.
  • MMD was the superior method, achieving the lowest $\beta$ across all pairs. For the challenging Simba $\to$ Eagle transfer, MMD reduced $\beta$ from 1.62 to 0.70 (a 57% reduction).
  • DANN showed mixed results; while it improved some transfers, it was unstable in others (e.g., slightly increasing $\beta$ when trained on EAGLE) due to the sensitivity of adversarial training to hyperparameters.
• Precision vs. Bias: HaloFlowDA (specifically MMD) achieved precision (posterior width) comparable to the standard Stellar-to-Halo Mass Relation (SHMR) method but with significantly lower bias and better calibration.
• Adaptive Uncertainty: The DA methods produced wider error bars in ambiguous cross-domain regions (where simulations disagree) and tighter bars where predictions were accurate, demonstrating adaptive uncertainty quantification.

5. Significance and Future Outlook

• Robustness to Physics Mismatch: The study demonstrates that domain adaptation can mitigate the "model misspecification" inherent in applying simulation-trained models to different physical regimes. This is a crucial step toward applying these tools to real observational data where the "true" generative process is unknown.
• MMD Superiority: The results suggest that kernel-based distribution matching (MMD) is more robust than adversarial training for this specific astrophysical problem, likely because it avoids the instability of the min-max game in adversarial learning.
• Path to Real Data: HaloFlowDA establishes a foundation for inferring halo masses from real HSC-SSP observations. The authors plan to validate these inferences against independent weak lensing and dynamical mass measurements.
• Generalizability: The findings extend to other information-maximizing compression methods in cosmology, highlighting that summaries trained on one simulation may fail out-of-domain without explicit adaptation.

In conclusion, HaloFlowDA successfully transforms a simulation-specific inference tool into a robust, generalizable framework, significantly reducing bias and improving calibration when bridging the gap between different cosmological simulations and, ultimately, real universe observations.