The Wrath of KAN: Enabling Fast, Accurate, and Transparent Emulation of the Global 21 cm Cosmology Signal

TL;DR

This paper introduces 21cmKAN, a fast and accurate emulator for the global 21 cm cosmology signal based on Kolmogorov-Arnold Networks, enabling rapid physical parameter estimation with high interpretability.

Contribution

The paper presents 21cmKAN, a novel KAN-based emulator that significantly improves training speed and maintains high accuracy over existing models like 21cmLSTM.

Findings

21cmKAN predicts signals in 3.7 ms on average.

Training time is reduced by a factor of 75 compared to 21cmLSTM.

Achieves unbiased posterior distributions in less than 30 minutes.

Abstract

Based on the Kolmogorov-Arnold Network (KAN), we present a novel emulator of the global 21 cm cosmology signal, $\texttt{21cmKAN}$, that provides extremely fast training speed while achieving nearly equivalent accuracy to the most accurate emulator to date, $\texttt{21cmLSTM}$. The combination of enhanced speed and accuracy facilitated by $\texttt{21cmKAN}$ enables rapid and highly accurate physical parameter estimation analyses of multiple 21 cm models, which is needed to fully characterize the complex feature space across models and produce robust constraints on the early universe. Rather than using static functions to model complex relationships like traditional fully-connected neural networks do, KANs learn expressive transformations that can perform significantly better for low-dimensional physical problems. $\texttt{21cmKAN}$ predicts a given signal for two well-known models in the community in 3.7 ms on average and trains about 75 times faster than $\texttt{21cmLSTM}$, when utilizing the same typical GPU. $\texttt{21cmKAN}$ is able to achieve these speeds because of its learnable, data-driven transformations and its relatively small number of trainable parameters compared to a memory-based emulator. We show that $\texttt{21cmKAN}$ required less than 30 minutes to train and fit these simulated signals and obtain unbiased posterior distributions. We find that the transparent architecture of $\texttt{21cmKAN}$ allows us to conveniently interpret and further validate its emulation results in terms of the sensitivity of the 21 cm signal to each physical parameter. This work demonstrates the effectiveness of KANs and their ability to more quickly and accurately mimic expensive physical simulations in comparison to other types of neural networks.