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

TL;DR

The paper tackles boundary pathology in simulation-based inference caused by sharp Uniform-prior sampling, which impairs neural posterior estimators near parameter-space edges. It introduces the Tailed-Uniform proposal, a hybrid density with Gaussian tails beyond prior bounds, providing smooth density transitions with minimal tuning. Through a Gaussian Linear toy and cosmological matter power-spectrum inference, the method consistently improves boundary fidelity and maintains strong performance across dimensions, often with far fewer simulations than Uniform requires. The approach is architecture-agnostic, computationally cheap, and publicly available, offering a practical upgrade to SBI workflows for robust edge inference.

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.

Figures (11)

• Figure 1: The standard Uniform distribution (blue) has constant probability density between -1 and 1 with zero probability outside this range. The Tailed-Uniform distribution (magenta) maintains a uniform density in the same central region $[-1, 1]$ but extends beyond these boundaries with Gaussian tails characterized by standard deviation $\sigma = 0.2$.
• Figure 2: Corner plots comparing posterior estimation performance for the boundary test case. The Tailed-Uniform (green) demonstrates superior boundary behavior compared to the Uniform (blue), closely matching the MCMC reference (yellow). The red dashed line indicates the true parameter value.
• Figure 3: Spatial distribution of C2ST performance across the parameter space, with blue regions indicating poor distributional matching (C2ST $\ll 0.5$) and red/orange regions indicating good performance (C2ST $\approx 0.5$) Top: Uniform versus analytical reference, showing systematic boundary degradation with glaring blue regions near parameter space edges. Bottom: Tailed-Uniform versus reference, demonstrating consistent performance across the entire parameter space.
• Figure 4: C2ST performance degradation as a function of distance from parameter space center. The blue curve reveals systematic deterioration of Uniform near boundaries, while the green curve demonstrates that Tailed-Uniform maintains consistent performance across all radii. The gray curve quantifies the increasing divergence between methods, with boundary regions showing substantial differences in posterior approximation quality. Values closer to 0.5 indicate better distributional matching. The error bars represent the 16th-84th percentile range of values at a fixed radius from the prior center.
• Figure 5: C2ST performance versus distance from center, stratified by tail width $\sigma$. The Uniform baseline (blue) exhibits systematic boundary degradation. The Tailed-Uniform's with varying tail widths (shown in green; lighter shades for smaller $\sigma$ values) exhibit robustness across a broad range of tail widths. Notably, performance becomes quite good for $\sigma \geq 0.2$ The error bars represent the 16th-84th percentile range of values at a fixed radius from the prior center.