Cross-simulator transfer with foundation model summaries: Towards robust SKA-era reionization inference

Yannic Pietschke; Caroline Heneka; Ayodele Ore; Romain Meriot

astro-ph.COAug 26, 2026

Cross-simulator transfer with foundation model summaries: Towards robust SKA-era reionization inference

Authors: Yannic Pietschke, Caroline Heneka, Ayodele Ore, Romain Meriot

Organizations: Institut für Theoretische Physik, Universität Heidelberg, Philosophenweg 16, 69120 Heidelberg, Germany · Max Planck Institute for Radioastronomy, Auf dem Hügel 69, Bonn, 53121, Germany

Abstract

Simulation-based inference (SBI) for parameter estimation is vulnerable to model misspecification: neural summaries and density estimators trained on a specific forward model typically fail when applied to data drawn from another model, or from real observations, and no training simulator can capture the full observational pipeline of a real measurement exactly. We show that a self-supervised Vision Transformer (ViT), pretrained label-free on a fast approximate simulator, produces transferable data summaries that generalize across simulators. Without retraining, it can be reused as a frozen encoder to infer astrophysical parameters from a completely different simulator that resolves the radiative transfer explicitly, on which it has never seen either data or parameters. As a concrete use case in 21cm cosmology, SKATR, a ViT pretrained with a Joint Embedding Predictive Architecture (JEPA), serves as a foundation model for reionization inference from upcoming SKA measurements: SKATR is pretrained once on 67k low-cost, noiseless semi-numerical 21cmFAST lightcones, then frozen and applied to hydrodynamical Loreli II lightcones, where a lightweight conditional flow matching head infers five astrophysical parameters; the encoder is never shown Loreli data, its parameters, or any noise. In our comparison, SKATR yields the most precise and best-calibrated posteriors across all five parameters, matching the accuracy of the fully-supervised in-domain baseline while requiring 2.6x fewer radiative-transfer simulations. Under realistic SKA AA* noise, only SKATR remains simultaneously accurate, informative, and calibrated, outperforming even a supervised baseline retrained from scratch on noisy data. Self-supervised pretraining on computationally efficient semi-numerical simulations is therefore a viable route to calibrated, simulator- and noise-agnostic reionization inference for the SKA-era.

Figures & tables

Parameter	Prior
$\log_{10}f_{\star,10}$	[-3, 0]
$\alpha_{\star}$	[-0.3, 0.9]
$\log_{10}f_{\mathrm{esc,10}}$	[-3, 0]
$\alpha_{\mathrm{esc}}$	[-0.8, 0.5]
$\log_{10}L_{\mathrm{X}}$ [erg s -1 ]	[38, 42]
$E_{0}$ [eV]	[100, 1500]

Table 1: Varied astrophysical parameters of the source (21cmFAST) dataset and their flat prior ranges.

Parameter	Meaning	Grid
$\tau$ [10 Myr]	star-formation timescale	$[700,\,10\,500]$
$M_{\mathrm{min}}$ [ $\mathrm{M_{\odot}}$ ]	min. star-forming halo mass	$[10^{8},\,4\times 10^{9}]$
$f_{X}$	X-ray production efficiency	$[0.1,\,10]$
$r_{\mathrm{H}}$	hard-to-soft X-ray ratio	$[0,\,1]$
$f_{\mathrm{esc,post}}$	late-time escape fraction	$\{0.05,\,0.275,\,0.5\}$

Table 2: Varied astrophysical parameters of the target (Loreli II) dataset and the values they take on the database grid ( Meriot et al., 2025 ) . 2 2 2 Notes.

M_{\mathrm{min}}

sampled jointly with

\tau

as 42 pairs;

f_{X}

sampled logarithmically over 13 values;

r_{\mathrm{H}}

sampled linearly over 6 values.

Figure 1: The three inference pipelines compared in this work. They share an identical downstream protocol: the frozen encoder maps a Loreli II lightcone to a 360-dimensional summary, which an ensemble of five conditional flow matching heads converts into a posterior over the Loreli astrophysical parameters. The pipelines differ only in how the encoder is pretrained. SKATR is pretrained self-supervised on 67,325 semi-numerical 21cmFAST lightcones with no parameter labels; ViT(21cmFAST) uses the same low-cost data but performs supervised regression on the six 21cmFAST parameters; ViT(Loreli) is trained supervised directly on $N$ expensive Loreli II lightcones, retrained for every training budget of $N\in\{100,500,1000,2500,5000,6517\}$ lightcones.

Panel of Figure 2 — Figure 2: Full budget of labeled Loreli hydrodynamical and RT lightcones ( $N=6517$ ), noiseless case. (a) Posterior corner plot for a representative test lightcone (truth in black), combined over the ensemble posteriors; the SKATR contours are the tightest and show no sign of bias. (b) Joint TARP coverage averaged over the ensemble and test set; SKATR is marginally underconfident, while the two supervised baselines carry a mix of biases and over-/under-confidence. The per-parameter breakdown becomes clear via the SBC ranks of Appendix A .

	$R^{2}$ ( $\uparrow$ )			Rel. unc. ( $\downarrow$ )
Parameter	ViT(Loreli)	ViT(21cmFAST)	SKATR	ViT(Loreli)	ViT(21cmFAST)	SKATR
$f_{X}$	0.97	0.97	0.99	0.12	0.11	0.06
$\tau$	0.95	0.94	0.99	0.14	0.18	0.08
$r_{\mathrm{H}}$	0.92	0.89	0.97	0.14	0.25	0.14
$\log M_{\mathrm{min}}$	0.98	0.98	1.00	0.11	0.13	0.06
$f_{\mathrm{esc,post}}$	0.96	0.91	0.97	0.12	0.20	0.09

Table 3: Per-parameter coefficient of determination

R^{2}

(

\uparrow

higher is better) and relative posterior width (defined as mean posterior standard deviation divided by the standard deviation of the label priors;

\downarrow

lower is more informative) at the full Loreli label budget (

N=6517

), computed on the combined ensemble posteriors. Best value per row is marked in bold.

Figure 3: Scaling with the number of Loreli training simulations

N

: accuracy (left, mean

R^{2}

) and posterior informativeness (right, mean relative uncertainty; lower is more informative). Errorbars show the spread over the five-member CFM ensembles; stars mark the combined ensemble. SKATR leads at every budget

N

of Loreli lightcones shown and exceeds the full-budget supervised (trained on

\sim

6.5k simulations) baseline at

N=2.5

Panel of Figure 4 — Figure 4: Inference under full Loreli budget of training lightcones ( $N=6517$ ) with SKA AA ∗ noise. (a) Posterior corner plot for a representative test lightcone (truth in black), combined over the $M=5$ CFM ensemble. SKATR is tightest with no signs of bias; the in-domain noise-trained ViT(Loreli) shows slightly broader constraints; the zero-shot ViT(21cmFAST) is broad and sits off the truth on several parameters. (b) Joint TARP coverage averaged over the ensemble and test set. SKATR tracks the diagonal (well-calibrated), the noise-trained ViT(Loreli) bows below it (overconfident and bias), and the zero-shot ViT(21cmFAST) lies near the diagonal because its posteriors are so broad and uninformative posteriors trivially achieve good coverage.

	$R^{2}$ ( $\uparrow$ )			Rel. unc. ( $\downarrow$ )
Parameter	ViT(Loreli, noise-trained)	ViT(21cmFAST)	SKATR	ViT(Loreli, noise-trained)	ViT(21cmFAST)	SKATR
$f_{X}$	0.94	0.85	0.96	0.26	0.27	0.11
$\tau$	0.88	0.66	0.92	0.21	0.44	0.19
$r_{\mathrm{H}}$	0.89	0.70	0.92	0.20	0.44	0.19
$\log M_{\mathrm{min}}$	0.94	0.80	0.97	0.19	0.37	0.14
$f_{\mathrm{esc,post}}$	0.87	0.69	0.91	0.22	0.45	0.20

Table 4: Per-parameter coefficient of determination

R^{2}

(

\uparrow

) and relative posterior width (

\downarrow

) at the full Loreli budget of training lightcones (

N=6517

) under SKA AA ∗ noise, computed on the ensemble posteriors. Best value per row in bold.

Figure 5: Scaling with the number of Loreli training simulations

N

under SKA AA ∗ noise: accuracy (left, mean

R^{2}

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Per-parameter rank statistics (SBC) at the full Loreli budget (

N=6517

), noiseless. SKATR and ViT(21cmFAST) are flat on four of the five parameters and mildly underconfident (

\cap

-shaped) on

f_{\mathrm{esc,post}}

; ViT(Loreli) shares this

f_{\mathrm{esc,post}}

underconfidence and is in addition visibly overconfident (

\cup

-shaped) on

r_{\mathrm{H}}

Figure 7: As Fig. 6 but under SKA AA ∗ noise. SKATR is essentially flat on every parameter (only a small hint of overconfidence on

r_{\mathrm{H}}

); the noise-trained ViT(Loreli) is

\cup

-shaped on several parameters, matching the joint-TARP overconfidence of Sec. 5.1 ; the zero-shot ViT(21cmFAST) shows indications of under- or overconfidence on several parameters.

Panel of Figure 8 — Figure 8: R+R augmentation on the supervised ViT(Loreli) baseline, noiseless. (a) Accuracy and informativeness versus the Loreli training budget $N$ for SKATR (no augmentation), ViT(Loreli) (no augmentation, the baseline from Sec. 3.3 ), and the R+R-augmented ViT(Loreli). The augmented baseline closes the accuracy and informativeness gap to SKATR by the full budget, and for known astrophysics and noise settings. (b) Joint TARP coverage at $N=6517$ for the same three pipelines: the R+R-augmented ViT(Loreli) now shows the same slight underconfidence as SKATR.

Figure 9: Per-parameter rank statistics (SBC) at the full Loreli budget (

N=6517

), noiseless, comparing SKATR (no augmentation), ViT(Loreli) (no augmentation), and the R+R-augmented ViT(Loreli). R+R augmentation removes the

\cup

-shape that the un-augmented ViT(Loreli) exhibits on

r_{\mathrm{H}}

, leaving the augmented baseline flat on every parameter except

f_{\mathrm{esc,post}}

, comparable to SKATR. The shared mild

\cap

-shape on

f_{\mathrm{esc,post}}

is more pronounced for ViT(Loreli) but persists across all three pipelines (see Sec. 4.1 ).

Explore similar work

STRADAViT: Self-Supervised Domain Adaptation of Vision Transformer Backbones for Radio Astronomy
Mar 31, 2026Andrea DeMarco, Ian Fenech Conti, Hayley Camilleri +2Domain-Adaptive Pretraining Self-Supervised Visual Representation Learning
Institute of Space Sciences and Astronomy, University of Malta, Msida, MSD2080, Malta · School of Physics and Astronomy, Royal Observatory, University of Edinburgh, Edinburgh, United Kingdom · Istituto Nazionale di Astrofisica (INAF), Osservatorio Astrofisico di Catania, Catania, Italy
Information-Preserving Domain Transfer with Unlabeled Data in Misspecified Simulation-Based Inference
May 7, 2026Joon Jang, Eunho Jeong, Kyu Sung Choi +1Synthetic-to-Real Domain Adaptation Simulation-Based Inference
Department of Biomedical Sciences, Seoul National University, Seoul, Republic of Korea · Department of Applied Bioengineering, Graduate School of Convergence Science and Technology, Seoul National University, Seoul, Republic of Korea · Department of Radiology, Seoul National University Hospital, Seoul, Republic of Korea +3
Towards foundation-style models for energy-frontier heterogeneous neutrino detectors via self-supervised pre-training
Apr 8, 2026Saúl Alonso-Monsalve, Fabio Cufino, Umut Kose +2Self-Supervised Pre-Training High-Energy Physics
IPA, ETH Zürich, Otto Stern Weg 5, Zurich, 8093, Switzerland.

Mar 31, 2026J/K move · Enter open · S save

STRADAViT: Self-Supervised Domain Adaptation of Vision Transformer Backbones for Radio Astronomy

Andrea DeMarco, Ian Fenech Conti, Hayley Camilleri, Ardiana Bushi+1

Institute of Space Sciences and Astronomy, University of Malta, Msida, MSD2080, Malta · School of Physics and Astronomy, Royal Observatory, University of Edinburgh, Edinburgh, United Kingdom · Istituto Nazionale di Astrofisica (INAF), Osservatorio Astrofisico di Catania, Catania, Italy

Domain-Adaptive Pretraining Self-Supervised Visual Representation Learning Vision Transformer Classification

Open paper AskPDF arXiv

Next-generation radio astronomy surveys are delivering millions of resolved sources, yet scalable morphology analysis remains difficult across heterogeneous telescopes and imaging pipelines. We present STRADAViT, a self-supervised continued-pretraining framework for learning transferable radio-astronomy encoders from Vision Transformer (ViT) backbones. It combines mixed-survey data curation, radio astronomy-aware training-view generation, and a ViT-MAE-initialized encoder family with optional register tokens. It supports reconstruction-only, contrastive-only, and two-stage branches. Our pretraining dataset comprises 512x512 radio astronomy cutouts drawn from four complementary sources (MeerKAT, ASKAP, LOFAR/LoTSS, and SKA SDC1 simulated data). We evaluate transfer with linear probing (LP) and fine-tuning (FT) on three morphology benchmarks spanning binary and multi-class settings (MiraBest, LoTSS DR2, and Radio Galaxy Zoo). An exploratory three-fold ablation grid guides selection of a register-based two-stage checkpoint using a fixed cross-dataset criterion. Across subsequent 15-seed paired downstream evaluations on fixed partitions, this checkpoint improves linear-probe Macro-F1 over its ViT-MAE initialization on all three benchmarks and improves fine-tuning on MiraBest and RGZ DR1, while LoTSS DR2 fine-tuning declines; all six differences remain statistically supported after Holm correction. A parallel DINOv2 experiment yields mixed adaptation effects: the procedure transfers, but the benefit is not uniform. STRADAViT thus improves frozen ViT representations while retaining clear dataset-dependent limitations and remaining below task-specialized methods on standard MiraBest classification.

Abstract

Explore similar work

STRADAViT: Self-Supervised Domain Adaptation of Vision Transformer Backbones for Radio Astronomy

Information-Preserving Domain Transfer with Unlabeled Data in Misspecified Simulation-Based Inference

Towards foundation-style models for energy-frontier heterogeneous neutrino detectors via self-supervised pre-training