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

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

Problem Statement
Simulation-Based Inference (SBI), particularly Neural Posterior Estimation (NPE), relies on training neural density estimators on simulated data pairs $(\theta, x)$ to approximate the posterior distribution $P(\theta|x)$. A common practice in large astrophysical campaigns (e.g., CAMELS, Quijote) is to generate training data by sampling parameters uniformly within a hard-truncated "box" defined by a Uniform prior.

The authors identify a critical structural pathology in this approach: the sharp discontinuity at the boundaries of a Uniform proposal distribution creates a region of sparse or non-existent training data near the edges. Because neural density estimators (such as flows) rely on gradient descent and continuous density transitions, they struggle to learn accurate approximations near these boundaries. Consequently, posterior estimates become unreliable or poorly constrained near the edges of the training box. This issue is exacerbated in high-dimensional spaces where the volume of the parameter space near boundaries dominates, and when the true posterior extends beyond the training proposal region.

Methodology: Tailed-Uniform
To address this, the authors propose Tailed-Uniform, a family of hybrid proposal distributions designed to pad the hard edges of the standard Uniform prior with decaying tails.

• Structure: The distribution maintains a flat density core (the original region of interest) but extends beyond the boundaries $[a, b]$ with a tail profile $h(\cdot)$. The tails are calibrated to ensure continuity at the boundary and to allocate a specific probability mass $A(\sigma)$ to the extrapolation region.
• Tail Shapes: The paper evaluates four specific tail profiles:
  1. Gaussian-tailed: Smooth, unbounded decay.
  2. Exponential-tailed: Smooth, unbounded decay.
  3. Linear-tailed: Compactly supported, linear decay.
  4. Flat-tailed: Compactly supported, constant density extension (a brute-force widening of the box).
• Training Mechanism: During training, the network minimizes the NPE loss function using importance weighting ($P(\theta)/\tilde{P}(\theta)$). This corrects for the mismatch between the proposal prior $\tilde{P}$ (Tailed-Uniform) and the assumed prior $P$ (e.g., a Uniform prior hard-truncated at the original box). This allows the network to learn the correct posterior under the assumed prior even though it was trained on data extending beyond that prior's support.
• Curse of Dimensionality: The authors analyze the trade-off in high dimensions. As dimensionality $d$ increases, the probability of a sample falling in the tail region grows exponentially ($1 - B(\sigma)^d$). They derive a heuristic for the maximum tail width $\sigma_{max}$ that keeps the tail allocation below a desired fraction of the total simulation budget.

Key Contributions and Experiments
The paper validates Tailed-Uniform through synthetic benchmarks and a cosmological application, comparing it against a standard Uniform baseline.

1. Toy Problems (SBI Benchmark Suite):

   • Tasks: Two tasks were used: a "Gaussian Linear" task (Gaussian assumed prior, allowing posterior leakage) and a "Gaussian Linear Uniform" task (Uniform assumed prior, hard-truncated).
   • Results: In both tasks, Tailed-Uniform variants consistently outperformed the Uniform baseline, particularly near the boundaries.
     • For the Gaussian prior, Uniform-trained networks failed to capture the posterior leakage beyond the box, while Tailed-Uniform networks successfully extrapolated.
     • For the Uniform prior (where the true posterior is zero outside the box), Tailed-Uniform still provided a consistent, albeit smaller, advantage. The extra samples outside the boundary helped constrain the network's behavior as it approached the edge, preventing the "diffuse" artifacts seen in Uniform-trained networks.
   • Metric: Performance was measured using the Classifier Two-Sample Test (C2ST), where scores closer to 0.5 indicate indistinguishability from the reference posterior. Tailed-Uniform maintained high scores (near 0.5) near boundaries, whereas Uniform scores dropped significantly (indicating poor performance).

2. Sensitivity Analysis:

   • Dimensionality: The advantage of Tailed-Uniform persists up to $d=8$ even when the tail region consumes ~83% of the simulation budget. At $d=12$, the advantage diminishes as the interior coverage becomes too sparse, but the method remains superior near boundaries compared to Uniform.
   • Data Volume & Architecture: Increasing the simulation budget by orders of magnitude (up to 10,000 samples) or increasing network width (up to 8x) did not resolve the boundary degradation for Uniform-trained networks. This confirms the issue is structural (lack of support beyond the edge) rather than a lack of data or model capacity.

3. Cosmological Application:

   • Task: Inference of four cosmological parameters ($\Omega_m, h, n_s, A_s$) from the matter power spectrum using the syren-new emulator.
   • Results: At a test point near the edge of the training box, Tailed-Uniform networks produced posteriors that closely matched MCMC references, whereas the Uniform-trained network showed significant deviation and poor constraint.
   • Degeneracies: The method remained robust even in the presence of parameter degeneracies (e.g., between $\Omega_m$ and $h$), maintaining performance in both constrained and degenerate directions near the boundaries.

Significance and Claims
The paper claims that Tailed-Uniform offers a robust, low-cost solution to the boundary pathology in SBI. The primary significance lies in demonstrating that sampling outside the assumed prior boundaries is beneficial even when the prior is hard-truncated.

• Structural Remedy: The degradation of Uniform-trained NPEs near boundaries is a structural failure of the proposal's support, not a limitation of data volume or network expressivity.
• Efficiency: Allocating a fraction of the simulation budget to tail regions yields significantly better posterior quality near boundaries than simply increasing the number of samples within the original box.
• Scalability: The method is particularly valuable in high-dimensional problems where boundary effects dominate the parameter space volume.
• Smoothness: Smooth tail profiles (Gaussian, Exponential) are generally preferred over flat or linear ones, as they are less sensitive to the specific choice of tail width $\sigma$.

The authors conclude that "cutting corners" (using hard-truncated Uniform proposals) is detrimental to inference robustness, and that "padding" the training box with decaying tails is a necessary step for reliable SBI in complex, high-dimensional astrophysical problems.