Don't Cut Corners: How Training Outside the Prior Makes Simulation-Based Inference More Robust

Large astrophysical simulation campaigns often generate training data by sampling parameters across a Uniform prior box. Due to the proposal’s sharp edge, neural posterior estimators struggle to learn accurate approximations near the boundaries. We propose Tailed-Uniform, a family of hybrid proposal distributions for sampling training simulations for robust simulation-based inference. By padding the original hard-truncated training box with decaying tails, Tailed-Uniform-trained networks yield more accurate posteriors near and beyond the edges. We demonstrate these improvements on a family of tail shapes, including a widened Uniform box as a control. Our results suggest that additional simulations near the prior boundary better constrain the networks as it approaches the edge of the training box, even for Uniform assumed priors. We show these advantages on a toy problem and cosmological parameter inference from the matter power spectrum. These benefits increase in high dimensions, where boundaries dominate parameter space volume. All code is publicly available on Github.

Improvements to AI systems:

1. Adaptive training-data proposal distributions for neural posterior estimators (NPEs): Replace hard-uniform priors with Tailed-Uniform distributions (e.g., Gaussian or logistic tails) during simulation-based training. This prevents sharp boundary artifacts, improving posterior accuracy near and beyond the prior edges—critical for scientific inference where true parameters may lie outside the assumed box.

2. Boundary-aware sampling for high-dimensional inference: In high-dimensional parameter spaces (where volume concentrates near boundaries), use Tailed-Uniform to allocate more simulations to edge regions. This directly reduces estimation error in dimensions where uniform sampling becomes inefficient, enabling more reliable posteriors for cosmological or particle physics models.

3. Robustness to prior misspecification: By training on tailed distributions, the AI system learns smoother likelihood approximations that remain valid even if the true posterior has support outside the original training box. This makes the system more reliable for real-world data where the true parameters are unknown and may exceed initial bounds.

4. Controlled tail-shape selection: The system can now tune tail decay (e.g., exponential vs. power-law) based on the target problem’s dimensionality and boundary sensitivity, providing a flexible hyperparameter for optimizing inference accuracy without changing the underlying neural architecture.

What the improved AI system can do:

• Accurately infer cosmological parameters (e.g., from matter power spectrum) even when the true values lie near or outside the prior box, reducing biased confidence intervals.

• Scale to high-dimensional inference tasks (e.g., 10+ parameters) where uniform sampling fails, maintaining calibration and coverage without requiring exponentially more simulations.

• Provide reliable uncertainty quantification for scientific discovery, enabling robust decision-making in fields like astronomy, particle physics, and climate modeling where boundary effects are common.

• Automatically generate training data that is more informative per simulation, cutting computational costs for expensive forward models (e.g., N-body simulations) while improving posterior fidelity.

Large astrophysical simulation campaigns often generate training data by sampling parameters across a Uniform prior box. Due to the proposal's sharp edge, neural posterior estimators struggle to learn accurate approximations near the boundaries. We propose Tailed-Uniform, a family of hybrid proposal distributions for sampling training simulations for robust simulation-based inference. By padding the original hard-truncated training box with decaying tails, Tailed-Uniform-trained networks yield more accurate posteriors near and beyond the edges. We demonstrate these improvements on a family of tail shapes, including a widened Uniform box as a control. Our results suggest that additional simulations near the prior boundary better constrain the networks as it approaches the edge of the training box, even for Uniform assumed priors. We show these advantages on a toy problem and cosmological parameter inference from the matter power spectrum. These benefits increase in high dimensions, where boundaries dominate parameter space volume.