Robust Simulation Based Inference Through Robust Optimal Transport

Technical Summary: Robust Simulation Based Inference Through Robust Optimal Transport

Problem Statement
Simulation Based Inference (SBI) is a framework for statistical inference when the likelihood function of a parametric model $\{P_\theta : \theta \in \Theta\}$ is analytically intractable, but simulations from the model are possible. Standard SBI methods often assume the data generating distribution $P$ is exactly a member of the model ($P = P_{\theta^*}$). However, in practice, models are rarely exact, and data may be subject to contamination. Existing robust methods typically address either "Huber contamination" (mass contamination where a fraction $\epsilon$ of data is replaced by arbitrary noise) or geometric perturbations (where data points are shifted by a small distance). A significant gap exists in methods that can simultaneously handle both types of contamination, particularly within the SBI setting where likelihoods are unavailable. Furthermore, many robust SBI approaches rely on Maximum Mean Discrepancy (MMD) or Wasserstein distance, which suffer from sensitivity to hyperparameter selection (e.g., kernel bandwidth) or lack robustness to mass contamination, respectively.

Methodology
The authors propose Bootstrapped Minimum Robust Semi-constrained Wasserstein-2 (B-MRSW), a framework that integrates a new divergence measure with the regular bootstrap for uncertainty quantification.

1. The Divergence ($\ell_\lambda$): The core of the method is the $\lambda$-Robust Semi-constrained Wasserstein-2 ($\lambda$-RSW) divergence. For probability measures $P_1, P_2$, it is defined as:
   $$\ell_\lambda(P_1, P_2) := \inf_{Q \ll P_1} \left[ \frac{1}{\lambda} KL(Q, P_1) + W_2^2(P_2, Q) \right]$$
   where $KL$ is the Kullback-Leibler divergence and $W_2$ is the Wasserstein-2 distance. This formulation relaxes the hard marginal constraint of standard optimal transport by allowing the first marginal to be reweighted (via $Q$) at a cost penalized by $KL$. This is motivated by distributionally constrained Empirical Likelihood.

   • Robustness: The authors prove that for intermediate values of $\lambda$ (e.g., $\lambda=1$), minimizing $\ell_\lambda(P, P_\theta)$ yields a parameter estimate robust to the Geometric+Huber (G+H) contamination model, where $P = (1-\epsilon)Q + \epsilon F$ and $W_2(Q, P_{\theta^*}) \le \rho$.
   • Limits: As $\lambda \to 0$, $\ell_\lambda$ converges to $W_2^2$ (sensitive to Huber contamination). As $\lambda \to \infty$, it converges to the average nearest-neighbor distance (uninformative for variance estimation).

2. Algorithm for Semi-Discrete Approximation:
   Since $P_n$ (the empirical data) is discrete and $P_\theta$ is often continuous, the problem is semi-discrete. The authors derive a dual formulation (Theorem 3.1) converting the min-min problem into a concave maximization problem over a vector $g \in \mathbb{R}^n$:
   $$\ell_\lambda(P_n, P_\theta) = \sup_{g \in \mathbb{R}^n} \mathbb{E}_{X \sim P_\theta} \left[ \min_{j \in [n]} (\|X - Y_j\|^2 - g_j) - \frac{1}{\lambda} \log \left( \frac{1}{n} \sum_{t=1}^n e^{-\lambda g_t} \right) \right]$$
   They propose Algorithm 1, a stochastic sub-gradient ascent (SGA) algorithm to approximate this value, providing a convergence guarantee of $O(\epsilon)$ expected absolute error in $O(n^2/\epsilon^2)$ time.

3. Inference Pipeline:

   • $\lambda$ Selection: A data-driven heuristic selects $\lambda$ by observing an "elbow" in the plot of the Wasserstein-2 distance between the fitted model and the reweighted empirical measure across a range of $\lambda$ values.
   • Parameter Estimation: The minimum divergence estimator $\hat{\theta}_\lambda$ is found using the Covariance Matrix Adaptive-Evolution Strategy (CMA-ES) to optimize the non-convex objective $\ell_\lambda(P_n, P_\theta)$, evaluated via Algorithm 1.
   • Uncertainty Quantification: The regular bootstrap is applied on top of the minimum $\lambda$-RSW estimator. $M$ bootstrap datasets are generated, and the CMA-ES/SGA pipeline is run for each to produce a distribution of estimates, forming confidence intervals.

Key Contributions

1. Theoretical Robustness: The paper establishes that the $\lambda$-RSW divergence provides theoretical robustness under the joint Geometric+Huber contamination model, a setting not previously addressed by a single unified method in SBI.
2. Algorithmic Development: The authors provide the first algorithm with a convergence guarantee for computing the semi-discrete Robust Semi-constrained Optimal Transport, utilizing a stochastic sub-gradient ascent approach.
3. Data-Driven Tuning: A principled, data-driven heuristic for selecting the robustness parameter $\lambda$ is introduced, removing the need for arbitrary user specification.
4. Uncertainty Quantification: The framework integrates the regular bootstrap to provide confidence sets with proper asymptotic coverage, addressing a common failure mode in robust Bayesian and frequentist inference.

Results
The method is evaluated on two benchmarks:

1. Univariate Normal Distribution: Under G+H contamination (mixing a Student's $t$ or discretized Normal with a Dirac mass), B-MRSW successfully recovers the mean and variance. In contrast, standard Minimum Wasserstein-2 estimation fails under Huber contamination.
2. G-and-K Distribution: A complex, four-parameter distribution used as a standard SBI benchmark. The authors compare B-MRSW against the NPL-MMD method (Dellaporta et al., 2022).
   • Sensitivity: NPL-MMD performance is highly sensitive to the Gaussian kernel bandwidth; small changes lead to massive errors in skewness estimation and poor coverage.
   • Robustness: B-MRSW maintains low Mean Squared Error (MSE) and high coverage rates across a wide range of $\lambda$ values selected by the heuristic.
   • Efficiency: While B-MRSW requires more simulator calls than NPL-MMD due to the use of CMA-ES for global optimization (vs. NPL-MMD's local gradient descent), the authors argue this is necessary to avoid local minima and ensure robustness.

Significance and Claims
The paper claims that B-MRSW offers a robust alternative for SBI in scenarios where models are misspecified and data contains both geometric and mass contamination. The authors emphasize that unlike existing methods which rely on careful, often difficult-to-tune hyperparameters (like kernel bandwidth in MMD), their approach uses a data-driven strategy for the robustness parameter. They note that while the method is computationally intensive due to the bootstrap and global optimization requirements, it is parallelizable. The work does not claim to solve all SBI problems but specifically targets the gap between geometric robustness and Huber robustness in likelihood-free settings, providing a tool that targets the true parameter (up to identifiability) while providing valid confidence sets. Future work is suggested to involve training non-parametric surrogates for the simulator to reduce computational cost.