Simulation-Based Inference for Cluster Cosmology with Set-Based Neural Network Architectures

Technical Summary: Simulation-Based Inference for Cluster Cosmology with Set-Based Neural Network Architectures

Problem Statement
The upcoming SRG/eROSITA All-Sky Survey provides an unprecedented statistical power through large catalogs of X-ray-selected galaxy clusters. Extracting cosmological constraints from these datasets requires robust statistical frameworks that accurately link observable quantities to underlying cosmological parameters. Traditional analyses rely on Markov Chain Monte Carlo (MCMC) methods, which typically involve numerical computation of expected cluster number densities via likelihood functions. These approaches face significant limitations: they often require compressing data into binned summary statistics (leading to information loss), struggle with high-dimensional model complexity due to the need to integrate out numerous nuisance parameters, and suffer from increased computational costs as sample sizes and ancillary data (e.g., weak lensing) grow. Furthermore, modeling complex survey selection functions, intrinsic scatter, and correlations between observables becomes increasingly difficult within standard likelihood frameworks.

Methodology
The authors develop and implement a Simulation-Based Inference (SBI) framework designed to bypass the need for an explicit likelihood function. The method operates in two distinct stages using neural networks trained on realistic mock catalogs calibrated to the eRASS1 survey:

1. Mock Generation Pipeline: A forward model generates synthetic galaxy cluster catalogs based on a subset of 11 parameters (5 cosmological, 5 X-ray scaling relation, and 1 correlation coefficient). The pipeline includes:

   • Sampling from the Halo Mass Function (HMF).
   • Applying scaling relations for X-ray count-rates and weak-lensing (WL) masses, including intrinsic correlated scatter.
   • Modeling the eRASS1 selection function (specifically an $L_{ext} > 10$ cut for high purity) and instrument scatter.
   • Generating tangential shear profiles (modeled after DES-like surveys) to incorporate weak lensing information.
   • The resulting mock catalogs contain variable numbers of clusters, with each cluster represented by an 8-dimensional vector (log count-rate, redshift, and six tangential shear values).

2. Neural Architecture:

   • Set-Based Embedding Network: To handle the variable number of clusters in each catalog without rigid binning, the authors employ a permutation-invariant "set-based" neural network (a variation of Deep Sets). This network processes individual clusters through an encoder ($\phi$) and aggregates the resulting feature vectors using a permutation-invariant aggregator. The aggregator computes statistics including the set cardinality ($\log N$), sum, mean, maximum, and a cross-feature correlation matrix. This produces a fixed-length, highly informative summary vector.
   • Neural Posterior Estimation: The summary vector is fed into a Masked Autoregressive Flow (MAF), a type of normalizing flow. The MAF learns the conditional posterior distribution $p(\theta | \hat{x})$ directly from the joint distribution of parameters and summary statistics generated by the mock pipeline. The training objective minimizes the negative log-likelihood of the true parameters given the summary statistics.

Key Contributions

• Variable-Length Input Handling: This work presents the first application of SBI to Bayesian cosmological inference for X-ray-selected clusters that does not require fixed input vector dimensions. By utilizing set-based neural networks, the method processes catalogs of varying sizes naturally, avoiding the information loss associated with traditional binning.
• Full Information Utilization: The framework utilizes the full cluster-level information content, including individual shear profiles and count-rates, rather than compressing data into pre-defined summary statistics.
• Robustness to Survey Effects: The method is tested against realistic survey effects, including selection biases and intrinsic scatter in scaling relations, demonstrating robustness through rigorous calibration tests.
• Scalability: The architecture is designed to be extensible to more complex forward models and additional observables without the computational explosion often seen in MCMC marginalization over nuisance parameters.

Results
The framework was validated using 200,000 mock catalogs. Key findings include:

• Parameter Recovery: The method successfully recovers input cosmologies within inferred uncertainties. For a mock sample size of $\sim 3,300$ clusters (matching the effective sample size of the $L_{ext} > 10$ eRASS1 subsample), the average constraints achieved are 11.5% on $\Omega_m$ and 4.4% on $\sigma_8$.
• Comparison to MCMC: Despite using a smaller subsample of clusters (3,259 vs. 5,259 in the G24 analysis) and omitting optical richness as a mass proxy, the SBI precision is comparable to traditional MCMC analyses based on larger samples. The SBI framework also demonstrated tighter recovery of the intrinsic scatter in the X-ray scaling relation compared to the MCMC approach.
• Calibration: The posterior estimator passed Simulation-Based Calibration (SBC) and Test of Accuracy with Random Points (TARP) diagnostics, confirming that the estimated posteriors are well-calibrated and unbiased, with only a slight tendency toward overconfidence.
• Performance: The summary network achieved a coefficient of determination ($R^2$) of approximately 0.93 on the validation set for the constrained parameters.

Significance
The paper positions this work as a methodological validation of SBI for next-generation large-scale structure analyses. The authors emphasize that while the current study focuses on mock data to validate the framework, the approach offers a flexible alternative to MCMC that can handle increasing model complexity and data volume without the prohibitive computational costs of high-dimensional integration. The framework is presented as readily extensible to future X-ray cluster surveys and multi-wavelength analyses (e.g., combining with Euclid or Roman Space Telescope data), offering a path toward efficient and accurate cosmological parameter estimation in the era of massive datasets. The authors note that applying this framework to actual observational data is deferred to future work, pending further adaptations to handle specific data characteristics like partial shear coverage.