Cross-Fitting for Neural Posterior Estimation

Neural Posterior Estimation (NPE) trains an amortized conditional density estimator to approximate posterior distributions in simulation-based inference. Each training pair is a single draw from the joint distribution of latent and observation, so posterior variance at an observation must be inferred from variation across pairs rather than within them. However, with a sufficiently flexible variational distribution, the NPE training objective is unbounded, and standard optimizers yield arbitrarily small variance. Regularizers like early stopping, dropout, and weight decay constrain network capacity without targeting the correct posterior variance. We propose Cross-Fitting NPE (CF-NPE), which separates mean estimation from higher-moment estimation via sample splitting. Concretely, CF-NPE estimates a conditional-mean ensemble by $K$-fold cross-fitting, then fits a conditional density to the resulting out-of-fold residuals. The conditional variance of an out-of-fold residual is bounded below by the true posterior variance, giving CF-NPE a well-defined population-level target that standard NPE lacks. Across eight synthetic benchmarks and a cosmological application, CF-NPE achieves higher held-out log-likelihood than standard NPE, with gains that grow with the latent dimension.