21cmKAN: Emulator for Global 21 cm Signal

• 21cmKAN is a machine-learning emulator that leverages Kolmogorov–Arnold Network layers to replace expensive simulations of the global 21 cm cosmology signal.
• It employs B-spline parameterized transformations to accurately emulate brightness temperature spectra with speed gains over traditional methods.
• The tool facilitates fast Bayesian inference with millisecond-scale evaluations, enabling efficient exploration of astrophysical and cosmological parameters.

21cmKAN is a machine-learning emulator of the global, sky-averaged 21 cm cosmology signal built on the Kolmogorov–Arnold Network (KAN) architecture. It is designed to replace expensive semi-numerical or semi-analytic forward models with a fast surrogate for Bayesian parameter inference, especially for upcoming global 21 cm measurements, including lunar far-side concepts where the radio environment is exceptionally quiet. In the published characterization of the method, 21cmKAN is described as fast, accurate, and relatively interpretable, with nearly equivalent emulation accuracy to 21cmLSTM, substantially faster training, millisecond-scale forward evaluation, and end-to-end training plus posterior inference in under 30 minutes for representative mock-data analyses (Jones et al., 15 Aug 2025).

1. Cosmological setting and computational role

The target observable is the global 21 cm signal, the sky-averaged differential brightness temperature of the redshifted hyperfine 21 cm line of neutral hydrogen relative to the radiation background. In the global-signal formulation, the observable is a one-dimensional function, usually written as $T_{21}(\nu)$ or $T_{21}(z)$, encoding the volume-averaged thermal and ionization history of the intergalactic medium across the Dark Ages, Cosmic Dawn, and the Epoch of Reionization. The rest-frame line lies at 1420 MHz and appears at low radio frequencies for the relevant redshifts, so the problem is observationally tied to low-frequency radiometry and theoretically tied to astrophysical and cosmological parameters that regulate Ly$\alpha$ coupling, X-ray heating, star formation, and ionization history (Jones et al., 31 Aug 2025).

In that setting, 21cmKAN addresses a computational bottleneck. Global-signal simulators such as 21cmFAST-like models, ARES, and related semi-analytic or semi-numerical frameworks are too expensive to evaluate $\sim 10^5$–$10^6$ times inside posterior-sampling pipelines. The central task is therefore to learn the mapping

$\boldsymbol{\theta} \mapsto \delta T_b(\nu),$

where $\boldsymbol{\theta}$ is a vector of cosmological and astrophysical parameters and the output is the full discretized spectrum. Once trained, the emulator substitutes millisecond-scale neural-network evaluation for repeated direct calls to the physical simulator. This is the practical basis for its use in nested sampling and related Bayesian workflows, particularly for future measurements from radio-quiet environments such as the lunar far side, including LuSEE-Night and FarView (Jones et al., 31 Aug 2025).

2. KAN architecture and mathematical formulation

21cmKAN is built from Kolmogorov–Arnold Network layers rather than conventional fully connected layers or recurrent cells. The architectural distinction is that standard multilayer perceptrons place scalar weights on edges and fixed activation functions on nodes, whereas KANs place learned univariate functions on edges and sum them at nodes. In a conventional network,

$y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$

where expressivity is accumulated through stacked linear maps and fixed nonlinearities. In a KAN, by contrast, the edge activation is itself learned and is parameterized as a smooth one-dimensional function, implemented in 21cmKAN with B-spline expansions (Jones et al., 15 Aug 2025).

The paper formulates the edge activation as

$\phi(x) \approx \sum_i c_i B_i(x),$

where $B_i(x)$ are local B-spline basis functions and $T_{21}(z)$0 are trainable coefficients. For a layer with input dimensionality $T_{21}(z)$1 and output dimensionality $T_{21}(z)$2, each edge from input node $T_{21}(z)$3 to output node $T_{21}(z)$4 learns its own $T_{21}(z)$5, and the node value in the next layer is the sum of incoming transformed quantities. This realizes the KAN principle of learning multivariate structure through compositions and sums of univariate functions, consistent with the Kolmogorov–Arnold representation theorem (Jones et al., 15 Aug 2025).

In the global 21 cm setting, the architecture is tailored to low-dimensional physical parameter spaces and smooth one-dimensional outputs. For the implementations reported in the KAN paper, the input layer contains 7 nodes for 21cmGEM or 8 nodes for ARES, the hidden stack consists of three KAN layers with 44, 44, and 71 nodes, and the output layer contains 451 or 449 nodes corresponding to brightness-temperature samples across the redshift grid. The splines are cubic, with 7 grid intervals. The resulting emulator is fully feed-forward, unlike 21cmLSTM, which treats the signal as a sequence and relies on recurrent gating and backpropagation through time (Jones et al., 15 Aug 2025).

3. Training data, preprocessing, and optimization

The reported implementations use two public synthetic libraries of global 21 cm signals. The first is 21cmGEM, a seminumerical model similar to 21cmFAST with Visbal/Fialkov-style physics. The second is ARES, a semi-analytic code that solves radiative transfer in 1D. For 21cmGEM, the training-plus-validation set contains 27,292 signals with a 90\%/10\% split and the test set contains 1,704 signals, sampled on $T_{21}(z)$6–$T_{21}(z)$7 with $T_{21}(z)$8. For ARES, the training-plus-validation set contains 26,552 signals and the test set again contains 1,704 signals on $T_{21}(z)$9–$\alpha$0 with $\alpha$1 (Jones et al., 15 Aug 2025).

The 21cmGEM parameterization includes 7 parameters, with examples given as $\alpha$2, $\alpha$3, $\alpha$4, $\alpha$5, $\alpha$6, $\alpha$7, and $\alpha$8. The ARES parameterization includes 8 parameters, with examples given as $\alpha$9, $\sim 10^5$0, $\sim 10^5$1, $\sim 10^5$2, $\sim 10^5$3, $\sim 10^5$4, $\sim 10^5$5, and $\sim 10^5$6. Physical constraints on the late-time neutral fraction are imposed during dataset generation: for example, $\sim 10^5$7 at $\sim 10^5$8 for 21cmGEM and $\sim 10^5$9 at $10^6$0 for ARES, so the training data already encode realistic reionization histories (Jones et al., 15 Aug 2025).

Preprocessing follows the published training setup closely. Parameters uniform in log-space are log-transformed, signals are flipped to train from high-$10^6$1 to low-$10^6$2, and global min–max normalization is applied to parameters and signal amplitudes using

$10^6$3

The loss is mean squared error between true and emulated normalized signals,

$10^6$4

Optimization uses a PyTorch backend with the Efficient-KAN implementation, Adam with learning rate $10^6$5 and weight decay $10^6$6, and a batch size of 100 in the final default configuration. Hyperparameter search is reported as having trained 1,000 KANs in parallel on an NVIDIA GH200 with Ray Tune and Optuna, exploring hidden-layer sizes, spline grid intervals, spline order, and batch size; three hidden layers gave the best trade-off between accuracy and complexity (Jones et al., 15 Aug 2025).

4. Emulation accuracy, speed, and Bayesian inference

Accuracy is quantified with the relative RMS error

$10^6$7

On 21cmGEM, averaged over 20 trials, 21cmKAN achieves a mean relative error of $10^6$8, a median error of $10^6$9, and a maximum error of $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$0, corresponding to an absolute mean error of approximately $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$1 mK. On ARES, the reported averages are $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$2 mean, $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$3 median, and $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$4 maximum, with an absolute mean error of approximately $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$5 mK (Jones et al., 15 Aug 2025).

The comparative results reported for the 21cmGEM benchmark are as follows.

Emulator	Mean error (%)	Max error (%)	Speed (ms)
21cmKAN	0.25	1.09	3.7
21cmLSTM	0.22	0.82	46
21cmVAE	0.35	1.84	41.4
globalemu	1.12	6.32	1.3
21cmGEM	1.59	10.55	160

These numbers define the standard positioning of 21cmKAN within the emulator landscape. It is nearly as accurate as 21cmLSTM, more accurate than 21cmVAE, and substantially more accurate than earlier polynomial or spline emulators, while remaining markedly faster than the recurrent alternative (Jones et al., 15 Aug 2025).

The speed claims are similarly explicit. Using a single NVIDIA A100 GPU and 10 CPU cores, training time is reported as approximately 10 minutes for 21cmGEM over 400 epochs and approximately 19.2 minutes for ARES over 800 epochs. On the same hardware, 21cmKAN trains about 75 times faster than 21cmLSTM and approximately 3–4 times faster than globalemu under the stated comparison setup. Mean forward evaluation time is approximately 3.7 ms per signal on the A100, and the paper further states that, when used side-by-side with 21cmLSTM in identical nested-sampling setups, 21cmKAN evaluates about 5 times faster (Jones et al., 15 Aug 2025).

The emulator is intended for direct insertion into Bayesian inference. The published demonstration uses MultiNest with 1,200 live points, sampling efficiency 0.8, and evidence tolerance 0.1. The likelihood is Gaussian in the residuals between mock data and emulator output. Three mock signals are selected from the 21cmGEM test set and one from the ARES test set, with Gaussian white noise added at levels $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$6, 10, and 25 mK, and in the ARES case also 50 mK. For representative 21cmGEM fits, the reported statistics include $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$7, $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$8, $\boldsymbol{\theta} \mapsto \delta T_b(\nu),$9 ms, and $\boldsymbol{\theta}$0 mK at $\boldsymbol{\theta}$1 mK, and $\boldsymbol{\theta}$2, $\boldsymbol{\theta}$3, $\boldsymbol{\theta}$4 ms, and $\boldsymbol{\theta}$5 mK at $\boldsymbol{\theta}$6 mK. For ARES with $\boldsymbol{\theta}$7 mK, the reported values are $\boldsymbol{\theta}$8, $\boldsymbol{\theta}$9, and $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$0 ms (Jones et al., 15 Aug 2025).

The same study reports posterior mean relative RMS errors with respect to the true signal of 0.91\% at $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$1 mK, 1.42\% at 10 mK, and 4.15\% at 25 mK for 21cmGEM, and 0.79\%, 1.38\%, and 2.93\% for ARES at 5, 10, and 25 mK respectively. The key inference claim is that 21cmKAN yields unbiased posterior distributions: for the examples presented, true parameter values are recovered within stated credible intervals, and end-to-end training and fitting can be completed in less than 30 minutes (Jones et al., 15 Aug 2025).

5. Interpretability, sensitivity structure, and comparison with related emulators

A distinctive aspect of 21cmKAN is that the learned transformations are inspectable. The published analysis visualizes first-layer activations by plotting the B-spline basis functions $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$2, the per-edge activations $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$3, and the aggregated node-wise activation $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$4 seen by the second layer. The color coding of activation curves is tied to the largest basis-coefficient magnitude, so the plots are read directly as sensitivity diagnostics in parameter space (Jones et al., 15 Aug 2025).

The reported empirical pattern is that parameters with large, strongly varying first-layer activations are also the ones with tighter posterior constraints and larger physical influence on the global signal. In 21cmGEM, $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$5, $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$6, and $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$7 exhibit large and highly variable activations and are correspondingly well constrained, whereas $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$8 shows small, relatively flat activations and weak posterior constraints. In ARES, $y_j = \sigma\Big(\sum_i w_{ji} x_i + b_j\Big),$9 and $\phi(x) \approx \sum_i c_i B_i(x),$0 show large activation magnitude and variability, while $\phi(x) \approx \sum_i c_i B_i(x),$1 and $\phi(x) \approx \sum_i c_i B_i(x),$2 are associated with weaker activations and broader posteriors. The paper also notes that the first three basis functions’ coefficients are near zero in all trials, so the lowest-order spline contributions vanish in practice (Jones et al., 15 Aug 2025).

This transparency is the principal contrast with 21cmLSTM. The LSTM approach treats the global signal as an ordered sequence in redshift or frequency and exploits strong temporal coherence through recurrent hidden and cell states. That design yields the lowest error currently reported for a global-signal emulator, and 21cmLSTM is described as the most accurate emulator in the community. Its disadvantages are sequential training, more parameters, heavier optimization, and substantially longer training times. 21cmKAN instead treats the task as static function approximation, with standard feed-forward backpropagation, lower memory pressure, and a smaller trainable model (Jones et al., 31 Aug 2025).

In the broader ML-for-21cm context, 21cmKAN belongs to the emulator branch rather than the direct 3D-inference branch. That distinction matters. Work on 3D tomographic lightcones has shown that shallow 3D CNNs are the best-performing models when the input is a volumetric $\phi(x) \approx \sum_i c_i B_i(x),$3 lightcone and the task is direct inference of parameters such as $\phi(x) \approx \sum_i c_i B_i(x),$4. There, “3D networks for 3D data work best,” because the data geometry is intrinsically volumetric (Heneka, 2023). By contrast, 21cmKAN addresses low-dimensional parameter-to-curve emulation, where KANs are especially well aligned with smooth functional composition.

6. Scientific significance, experimental relevance, and future directions

The immediate scientific role of 21cmKAN is to make full Bayesian analysis of global 21 cm measurements computationally practical. This matters because the global 21 cm signal is sensitive to first-galaxy astrophysics, heating histories, reionization timing, and beyond-$\phi(x) \approx \sum_i c_i B_i(x),$5CDM scenarios. It is also relevant to lunar far-side observing programs, where the absence of terrestrial radio-frequency interference and the uniquely quiet radio environment motivate end-to-end forward-modeling pipelines that can exploit clean data streams without making the inference stage itself a bottleneck (Jones et al., 31 Aug 2025).

A broader implication is that emulator speed changes the scale of feasible model exploration. The published discussion emphasizes that global 21 cm analyses will require multiple emulators trained on different underlying physical simulators and parameterizations. Because 21cmKAN can be retrained in minutes rather than many hours, emulator training ceases to be the dominant cost in a campaign that compares multiple astrophysical prescriptions, foreground treatments, or nuisance-parameter choices. This suggests a natural role for 21cmKAN in multi-model studies, joint analyses with high-$\phi(x) \approx \sum_i c_i B_i(x),$6 UV luminosity functions, and rapid iteration as experimental data products evolve (Jones et al., 15 Aug 2025).

The method also fits into a wider shift toward simulation-based inference in 21 cm cosmology. On the forward-model side, codes such as 21cmFirstCLASS aim to couple CLASS-based cosmological initial conditions to 21cmFAST-like semi-numerical evolution so that cosmology and astrophysics are treated consistently from recombination onward (Flitter et al., 2023). On the observational side, interferometric analyses continue to confront calibration, beam chromaticity, polarization leakage, and wedge contamination, all of which structure how 21 cm observables are constructed from visibilities and power spectra (Bernardi, 2019). 21cmKAN does not solve those measurement-systematics problems directly; its contribution is to the fast-surrogate layer of the inference stack.

The published future directions are correspondingly focused. The authors identify training on more complex observables such as 3D 21 cm power-spectrum summaries, combining global signal, power spectrum, and galaxy UV luminosity functions, extending KANs to temporal or light-cone representations, using KAN-based learned transformations for domain adaptation between simulators such as 21cmGEM and ARES, and embedding the emulator in fuller experimental pipelines that jointly model signal, foregrounds, and instrumental systematics. A plausible implication is that 21cmKAN is best understood not as a terminal architecture but as a specialized, low-dimensional emulator design within a larger 21 cm inference ecosystem (Jones et al., 15 Aug 2025).