Cross fitting improves neural posterior estimation accuracy and calibration
Cross-Fitting for Neural Posterior Estimation
Machine Learning
Summary
Estimating the probability of unknown things based on observed data can sometimes lead to models that are too confident and not reliable on new data. The authors found that training these models normally causes them to overfit, making them less trustworthy. They invented a method called Cross-Fitting NPE that separates the learning of the average prediction from the uncertainty part, by splitting the data and training multiple models. This method makes the models perform better on unseen data and be better calibrated, especially when data is limited or complex.
What this means in practice
- •For simulation engineers: Improve uncertainty quantification in simulation-based models when data samples are limited and parameter spaces are high dimensional.
- •For cosmology data teams: Increase accuracy and reliability of parameter inference from cosmological simulations through better posterior modeling.
- •For machine learning platform developers: Enhance neural density estimation tools in platforms by integrating cross-fitting to prevent overconfident posterior estimates.$Commercial implications: Enables building reliable probabilistic inference products for users needing trustworthy uncertainty estimates, differentiating offerings.
Authors
Jeffrey Regier
Abstract
Neural Posterior Estimation (NPE) fits an amortized conditional density estimator to a simulated training set of complete-data realizations. NPE's training objective, log-likelihood, favors posterior approximations that assign arbitrarily high density to the training points, and typical families of amortized conditional density estimators (e.g., normalizing flows, Gaussian encoders) are flexible enough to oblige. The resulting policy of extreme overconfidence translates poorly to held-out data. Early stopping and other regularizers can prevent this degeneracy, but they are limited to selecting posterior approximations along an unfavorable optimization trajectory set by in-sample residuals, where the posterior mean is underfit while the posterior variance is overfit. To avoid this trade-off, we propose Cross-Fitting NPE (CF-NPE), which uses sample splitting to separate conditional-mean estimation from the estimation of the posterior's higher moments. Using $K$-fold cross-fitting, CF-NPE fits $K$ posterior-mean regressors and then fits a conditional density estimator to the resulting out-of-fold residuals. In experiments with synthetic data and in a cosmological application, CF-NPE achieves higher held-out log-likelihood than standard NPE and typically improves calibration. CF-NPE is most useful in high-dimensional, simulation-limited settings.