Learning at the Edge: Tailed-Uniform Sampling for Robust Simulation-Based Inference

TL;DR

This paper proposes the Tailed-Uniform distribution for simulation-based inference, extending sampling beyond prior bounds with Gaussian tails to improve neural posterior estimation near boundaries, demonstrated on synthetic and cosmological tasks.

Contribution

Introduction of the Tailed-Uniform proposal distribution that enhances robustness and boundary performance in simulation-based inference with minimal hyperparameter tuning.

Findings

Tailed-Uniform improves boundary performance over standard uniform sampling.

The method is robust with tail widths of 10-30% of prior range.

Performance gains increase with higher parameter space dimensions.

Abstract

We introduce the \textsc{Tailed-Uniform} proposal distribution for generating training simulations in simulation-based inference. Instead of sampling parameters uniformly within bounded regions, we extend the distribution beyond prior boundaries with smooth Gaussian tails. This eliminates sharp discontinuities that cause neural posterior estimators to fail when the posterior distribution intersects or extends beyond the prior bounds. The method requires minimal hyperparameter tuning, with tail widths of 10--30\% of the prior width proving robust across problems. We demonstrate these benefits on a synthetic Gaussian linear task and cosmological parameter inference from the matter power spectrum. We also find that \tail-trained models outperform \textsc{Uniform} ones near the boundaries across various training set sizes and dimensions of the parameter space. This advantage grows in higher dimensions, where boundaries dominate parameter space volume. All code is publicly available on Github at https://github.com/chaipattira/tailed-uniform-sbi.