Attention-Based Neural Network Emulators for Multi-Probe Data Vectors Part II: Assessing Tension Metrics

TL;DR

This paper advances neural network emulators using transformer architectures to efficiently analyze large cosmological datasets, significantly improving accuracy and scalability for Bayesian inference in cosmology.

Contribution

It introduces a generalized transformer-based neural network architecture with enhanced versatility and a scalable training method, outperforming previous models in accuracy for cosmological data analysis.

Findings

Transformer-based emulators outperform previous neural network models in accuracy.

The new architecture scales efficiently to larger datasets and parameter spaces.

Emulators effectively calibrate tension metrics in cosmological inferences.

Abstract

The next generation of cosmological surveys is expected to generate unprecedented high-quality data, consequently increasing the already substantial computational costs of Bayesian statistical methods. This will pose a significant challenge to analyzing theoretical models of cosmology. Additionally, new mitigation techniques of baryonic effects, intrinsic alignment, and other systematic effects will inevitably introduce more parameters, slowing down the convergence of Bayesian analyses. In this scenario, machine-learning-based accelerators are a promising solution, capable of reducing the computational costs and execution time of such tools by order of thousands. Yet, they have not been able to provide accurate predictions over the wide prior ranges in parameter space adopted by Stage III/IV collaborations in studies employing real-space two-point correlation functions. This paper offers a leap in this direction by carefully investigating the modern transformer-based neural network (NN) architectures in realistic simulated Rubin Observatory year one cosmic shear $\Lambda$CDM inferences. Building on the framework introduced in Part I, we generalize the transformer block and incorporate additional layer types to develop a more versatile architecture. We present a scalable method to efficiently generate an extensive training dataset that significantly exceeds the scope of prior volumes considered in Part I, while still meeting strict accuracy standards. Through our meticulous architecture comparison and comprehensive hyperparameter optimization, we establish that the attention-based architecture performs an order of magnitude better in accuracy than widely adopted NN designs. Finally, we test and apply our emulators to calibrate tension metrics.