Direct Learning of Calibration-Aware Uncertainty for Neural PDE Surrogates
This addresses the need for reliable uncertainty estimation in scientific computing applications, but it is incremental as it builds on existing neural PDE surrogate methods.
The paper tackles the problem of obtaining calibrated uncertainty for neural PDE surrogates in data-limited regimes by proposing a cross-regularized uncertainty framework that learns uncertainty parameters during training, resulting in better calibration on held-out splits and uncertainty fields that concentrate in high-error regions.
Neural PDE surrogates are often deployed in data-limited or partially observed regimes where downstream decisions depend on calibrated uncertainty in addition to low prediction error. Existing approaches obtain uncertainty through ensemble replication, fixed stochastic noise such as dropout, or post hoc calibration. Cross-regularized uncertainty learns uncertainty parameters during training using gradients routed through a held-out regularization split. The predictor is optimized on the training split for fit, while low-dimensional uncertainty controls are optimized on the regularization split to reduce train-test mismatch, yielding regime-adaptive uncertainty without per-regime noise tuning. The framework can learn continuous noise levels at the output head, within hidden features, or within operator-specific components such as spectral modes. We instantiate the approach in Fourier Neural Operators and evaluate on APEBench sweeps over observed fraction and training-set size. Across these sweeps, the learned predictive distributions are better calibrated on held-out splits and the resulting uncertainty fields concentrate in high-error regions in one-step spatial diagnostics.