Efficient self-attention-based emulation of gravitational-wave spectrum with full parameter space exploration

Abstract Fast emulation of sequential data generated by numerical computations is vital for large-scale Bayesian analysis in cosmology and astrophysics. Despite the success of neural network models in such tasks, benchmarks typically fall short of the required sub-percent accuracy and lack generalizability to broader data sets of different features and physical origins. Here we introduce SageNet+, a physics-informed framework for emulating gravitational-wave spectrum, while our self-attention-based design is adaptable to generic sequential data. First-principles modeling of the stochastic gravitational-wave background from cosmic inflation takes tens of seconds per evaluation to solve differential integrations. Instead, our SageNet+ emulator can predict the present-day gravitational-wave spectrum in less than 10 ms for every set of the nine cosmological parameters with wide prior ranges, achieving a thousand-fold speed-up over numerical methods. Based on the Transformer architecture, SageNet+ combines learnable positional embeddings and adaptive sequence standardization to preserve spectral features. We further employ refined data preprocessing to maximize physical information and devise novel metrics to assess model performance. Trained on 180,000 spectra that explore the full parameter space, SageNet+ attains sub-percent accuracy for all test data. This enables real-time parameter inferences and joint analyses with external data. SageNet+ presents an optimal, systematic and generalizable framework for emulation tasks in physical sciences.