PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yannic Pietschke, Caroline Heneka, Ayodele Ore, Romain Meriot

arXiv:2608.26354v1astro-ph.COastro-ph.IMcs.LGphysics.data-an

TL;DR

SBI can fail under simulator and observational domain shifts, motivating transferable summaries that are not tied to one forward model. The paper pretrains SKATR label-free with JEPA on inexpensive 21cmFAST lightcones, freezes it, and applies it with a conditional flow-matching head to Loreli II data. SKATR matches supervised in-domain accuracy with 2.6× fewer radiative-transfer simulations and remains accurate, informative, and calibrated under realistic SKA AA* noise.

  • Problem

    SBI summaries and density estimators trained on one simulator typically fail on other simulators or real observations, where model misspecification threatens reliable inference.

  • Method

    SKATR is a JEPA-pretrained Vision Transformer trained label-free on noiseless 21cmFAST lightcones, frozen, and paired with a conditional flow-matching head for Loreli II parameter inference.

  • Results

    SKATR matches the fully supervised baseline while using 2.6× fewer radiative-transfer simulations and is the only compared pipeline simultaneously accurate, informative, and calibrated under realistic SKA noise.

  • Takeaways & Limitations

    Label-free pretraining on inexpensive semi-numerical simulations supports calibrated transfer across simulator and noise shifts for SKA-era reionization inference.

  • Takeaways & Limitations

    The noise model omits the full calibration chain, residual foreground-subtraction errors, and beam systematics, and robustness beyond Loreli II remains untested.

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses SBI’s vulnerability to simulator and observational domain shifts by testing self-supervised foundation-model pretraining for transferable 21cm summaries. SKATR is pretrained label-free on inexpensive 21cmFAST lightcones, frozen, and applied to hydrodynamical Loreli II data with realistic SKA noise.

  • Motivation: SBI neural summaries and density estimators commonly fail when data come from another simulator or real observations.Distinct parameterizations, physical approximations, and incomplete observational modeling create model misspecification.
  • Motivation: Self-supervised foundation models pretrain a label-free encoder on abundant low-cost data, then reuse its frozen representation with light task-specific supervision.Because the objective is not tied to one target-label set, representations can transfer to unseen parameters, noise models, and systematics.
  • 21cm application: 21cm cosmology is a stringent test because SKA measurements include thermal noise, ionospheric effects, foreground residuals, beam and calibration artifacts, while simulators use differing approximations.These shifts affect posterior accuracy, precision, and calibration.
  • Approach: SKATR tests whether JEPA-based, symmetry-agnostic pretraining yields summaries transferable from inexpensive semi-numerical lightcones to expensive hydrodynamical simulations.The encoder is pretrained once on noiseless data, frozen, and applied to Loreli II lightcones with and without realistic SKA-Low noise.
  • Evaluation: The study evaluates cross-simulator transfer and realistic-noise transfer using a lightweight conditional flow-matching inference head.The encoder sees neither Loreli data or parameters nor noise during pretraining.

2. Simulation methods

The study transfers representations from a large, inexpensive 21cmFAST source set to a smaller, computationally expensive Loreli II target set with explicit radiative transfer. It controls data leakage and adds a conservative SKA noise-and-foreground domain shift for evaluation.

  • 21cmFAST source data: The 21cmFAST source fields model 21cm brightness temperature across Cosmic Dawn and reionization, varying six astrophysical source parameters.The parameters include stellar and ionizing escape-fraction terms and X-ray-heating controls.
  • Loreli II target data: Loreli II varies five astrophysical parameters on a structured grid, including star-formation, escape-fraction, and X-ray-emission controls.The late-time escape fraction takes discrete values, while fX and rH are sampled over specified grids.
  • Computational motivation: A complete Loreli II lightcone requires 300–500 CPU-h, motivating transfer from a source suite that is approximately 180× cheaper overall.The source suite costs ∼1.7 × 10^4 CPU-h versus ∼3.1 × 10^6 CPU-h for the full Loreli II suite, without matched resolution.
  • Data control: The source lightcones are cropped to the common redshift interval z = 5–13.15 before training, and 1,150 Loreli II lightcones are held out for every budget.Training budgets draw from the remaining 6,517 target lightcones, preventing leakage across training and evaluation.
  • Noise model: The noise study corrupts Loreli II lightcones with SKA AA* thermal noise and a conservative foreground-wedge cut.The wedge extends 0.1 h Mpc^-1 beyond the horizon limit and uses the AA* configuration with 1080 h integration.

3. Network architecture and experimental setup

The architecture freezes a JEPA-pretrained ViT encoder that compresses each lightcone into a 360-dimensional summary, then uses conditional flow matching for posterior inference. Experiments isolate pretraining strategy while varying target-data budgets and averaging inference-head randomness.

  • SKATR encoder: SKATR tokenizes each 32 × 32 × 576 lightcone into 3,072 non-overlapping patches and projects them into 360-dimensional token embeddings.Learnable Fourier sin/cos positional encodings are added to the tokens.
  • SKATR encoder: JEPA trains SKATR to predict masked-region embeddings from visible context without per-sample labels or voxel-space reconstruction.A context encoder, fixed target encoder, and predictor transformer implement the embedding-space prediction task.
  • SKATR encoder: Masking uses four rectangular target regions at mixed spatial scales, with two covering 60–80% and two covering 10–20% of the spatial area.Each region spans the full redshift axis, so masking acts only in the sky plane.
  • SKATR encoder: Unlike canonical JEPA, the target encoder is initialized as a copy of the context encoder and then held fixed rather than updated by exponential moving average.Each target rectangle is predicted independently from its complementary context.
  • Transfer protocol: After ∼60 GPU-hours of pretraining, the encoder is frozen and mean-pools 3,072 token embeddings into a 360-dimensional Loreli summary.No Loreli data, target parameters, or noise are shown during pretraining.
  • Inference head: Conditional flow matching maps each frozen summary to a posterior over five Loreli parameters by transporting a Gaussian base distribution through a learned continuous-time ODE.Posterior samples are generated by integrating the learned velocity field from t = 0 to t = 1.
  • Experimental setup: All pipelines use the same ViT backbone family, frozen 360-dimensional summaries, and downstream CFM protocol; only encoder pretraining differs.The comparison includes self-supervised SKATR, supervised ViT(21cmFAST), and budget-specific supervised ViT(Loreli).
  • Experimental setup: Five independently initialized CFM heads are trained for every pipeline and target-lightcone budget, and evaluation uses accuracy, informativeness, and calibration.The combined posterior contains 2,500 samples per lightcone, while accuracy is assessed with per-parameter and mean R2.

4. Results I: Cross-simulator transfer

SKATR transfers a frozen, self-supervised representation from 21cmFAST to Loreli II, producing tighter, more accurate, and better-calibrated five-parameter posteriors than supervised baselines. Its advantages persist as labeled target simulations become scarce, yielding equal-or-better accuracy with substantially fewer expensive Loreli simulations.

  • Full-budget posteriors: At N = 6517, SKATR achieves the best mean R2 across all five parameters, reaching 0.98 while keeping contours tight and centered on truth.It has the smallest relative posterior width on four parameters and ties ViT(Loreli) at 0.14 on rH.
  • Full-budget posteriors: SKATR improves mean R2 from 0.94 for ViT(21cmFAST) to 0.98, with the largest gains on rH and fesc,post.The self-supervised objective, rather than pretraining data scale alone, is identified as producing a more effective embedding.
  • Calibration: SKATR is calibrated on four parameters and slightly conservative on one, whereas supervised baselines show bias or overconfidence.The shared underconfidence on fesc,post is associated with sparse labels and the continuous flow-matching posterior model.
  • Data scaling: At N = 500, SKATR reaches R2 = 0.90 versus 0.76 for ViT(21cmFAST) and 0.59 for ViT(Loreli), retaining the best calibration across budgets.SKATR remains ahead at every tested training budget and has the smallest run-to-run scatter across ensemble initializations.
  • Simulation budget efficiency: SKATR trained on N = 2500 Loreli simulations reaches R2 = 0.97, exceeding the fully supervised full-budget baseline at R2 = 0.96.This corresponds to a 2.6× reduction in target-simulator labels for equal or better accuracy.
  • Simulation budget efficiency: SKATR’s frozen encoder achieves comparable full-budget accuracy without target-simulator or noise exposure, while target-specific tuning can close the supervised accuracy gap only at full budget.Its data-efficiency advantage remains in the label-scarce regime and carries into realistic noisy mocks.

5. Results II: Cross-simulator transfer under realistic SKA noise

Under realistic SKA AA* noise, SKATR remains accurate, informative, and calibrated across the full and reduced Loreli training budgets. It outperforms zero-shot supervised transfer and the noise-trained in-domain baseline on the combined evaluation.

  • Posteriors and calibration at the full budget: The noise-trained ViT(Loreli) maintains accuracy but produces systematically too-narrow, overconfident posteriors.Its joint TARP curve falls below the diagonal and its SBC rank histograms are ∪-shaped for several parameters.
  • Posteriors and calibration at the full budget: SKATR is the only pipeline simultaneously accurate, informative, and calibrated under realistic SKA AA* noise.Its posteriors remain tight and centered on the truth, while its joint TARP coverage is close to the diagonal.
  • Posteriors and calibration at the full budget: Mean R2 is 0.94 for SKATR, compared with 0.91 for noise-trained ViT(Loreli) and 0.74 for zero-shot ViT(21cmFAST) at full budget.SKATR also has the smallest mean relative posterior width, 0.17, versus around 0.22 and 0.40 for the two baselines.
  • Data scaling and robustness: SKATR leads in mean R2 at every noisy Loreli budget, increasing from 0.79 at N = 500 to 0.94 at N = 6517.At N = 500, it achieves 0.79 versus 0.65 for the noise-trained ViT(Loreli) and 0.55 for zero-shot ViT(21cmFAST).
  • Data scaling and robustness: With identical architecture, pretraining data, and downstream protocol, SKATR’s mean R2 drops 4% under the input shift, versus 21% for ViT(21cmFAST).The comparison isolates the self-supervised versus supervised pretraining objective as the reported source of the robustness difference.

6. Discussion and conclusions

SKATR transfers from low-cost semi-numerical pretraining to hydrodynamical simulations and realistic SKA noise while remaining accurate, informative, and calibrated. The main boundaries are incomplete instrument-systematics coverage, testing only one target simulator family, and an unresolved explanation of why JEPA features transfer.

  • Results: SKATR reaches mean R2 = 0.98 versus 0.96 for supervised ViT(Loreli) on noiseless Loreli II data, with the tightest posterior across all five parameters.Both methods use the same budget of training simulations.
  • Results: Under realistic SKA AA∗ thermal noise and wedge avoidance, SKATR retains 96% of noiseless accuracy as mean R2 drops from 0.98 to 0.94.The same frozen encoder remains accurate despite not seeing the target simulator or noise during pretraining.
  • Results: 2.6 fewer target simulations are needed: at N = 2500 target lightcones, SKATR matches fully supervised ViT(Loreli) at N = 6517.The comparison reports this saving in target radiative-transfer simulations.
  • Interpretation: JEPA transfer is a working hypothesis: masked-voxel prediction may preserve physics-related information under simulator changes, noise, and foreground-wedge removal.The authors state that direct latent-space analysis is needed to confirm this mechanism.
  • Results: SKATR is the only compared pipeline simultaneously accurate, informative, and calibrated under the noisy joint simulator-and-noise shift.It outperforms an in-domain baseline retrained on noisy target data on accuracy, informativeness, and calibration.
  • Limitations: The noise model omits the full calibration chain, residual foreground-subtraction errors, and beam systematics.Extending agnosticism to these effects is identified as a next step.
  • Limitations: Robustness beyond the single Loreli II target simulator family remains untested.The paper also notes that baseline comparisons use a controlled shared architecture and training protocol.
  • Outlook: Direct latent-space analysis could test whether the 360-dimensional representation clusters by physical reionization regimes rather than simulator family or noise artifact.This is presented as the most interesting follow-up study.

Data availability

The work builds on publicly available SKATR and uses the publicly available EoRFlow framework for inference.

  • Data availability: SKATR is publicly available through its GitHub repository.The paper states that the work builds upon the SKA Transformer.
  • Data availability: Inference was performed with the publicly available EoRFlow framework.The framework is identified by its GitHub repository.

Appendix A: Marginal calibration: rank statistics

Per-parameter simulation-based calibration evaluates whether posterior ranks are uniformly distributed under the prior-simulator distribution. The appendix finds broadly good calibration for SKATR, especially under SKA noise.

  • Rank statistics: Simulation-based calibration records each true parameter’s rank among 2500 combined-ensemble posterior samples over 1150 held-out test lightcones.The rank is defined as the fraction of posterior samples below the truth.
  • Rank statistics: Uniform rank histograms indicate correctly calibrated marginal posteriors, while endpoint accumulation signals overconfident posteriors.The appendix uses histogram shape as a low-variance miscalibration diagnostic.
  • Results: In noiseless data, SKATR is well calibrated on four of five parameters and mildly underconfident on fesc,post.ViT(Loreli) is additionally overconfident on rH.
  • Results: Under SKA AA∗ noise, SKATR’s per-parameter ranks are essentially uniform across all five parameters.This supports the joint calibration result reported elsewhere in the paper.

Appendix B: Augmentation and the limits of supervised tuning

R+R augmentation substantially strengthens the supervised in-domain ViT baseline, but SKATR retains an advantage when target simulations are scarce and remains more reusable across settings.

  • Augmentation results: R+R augmentation closes the supervised ViT(Loreli) accuracy, informativeness, and calibration gap to SKATR at the full Loreli budget.The augmented baseline matches SKATR’s mean R2 within statistical scatter and roughly matches its joint-TARP calibration.
  • Data efficiency: +0.12 is SKATR’s combined-ensemble mean R2 lead at N = 500, compared with +0.04 at N = 1000 and +0.01 at N = 2500.The lead vanishes at the full Loreli budget, showing that augmentation can match SKATR when enough expensive target simulations are available.
  • Data efficiency: Sky-plane augmentation does not substitute for SKATR pretraining when target labels are scarce.SKATR reuses one frozen encoder across budgets, whereas the supervised baseline trains a separate target-specific encoder at each budget.
  • Limits of supervised tuning: R+R requires identifying the relevant sky-plane symmetry beforehand, while SKATR’s encoder remains parameter-set-agnostic and avoids target-specific backbone retraining.The supervised regression objective is tied to the five Loreli parameters, whereas those parameters are seen only by SKATR’s lightweight CFM head.
  • Noise robustness: Under SKA AA* noise, SKATR remains essentially flat across parameters, while the noise-trained ViT(Loreli) is overconfident on several parameters and zero-shot ViT(21cmFAST) shows mixed calibration indications.Only SKATR avoids visibly broad over- or underconfidence in this comparison.
  • Calibration: R+R removes the un-augmented ViT(Loreli) overconfidence on rH, leaving rank statistics flat except for shared mild underconfidence on fesc,post.SKATR and the augmented baseline are similarly calibrated on the remaining parameters.
Loading 2608.26354v1…