Direct Learning of Calibration-Aware Uncertainty for Neural PDE Surrogates

TL;DR

This paper introduces a novel training framework for neural PDE surrogates that learns calibrated, regime-adaptive uncertainty during training, improving the reliability of uncertainty estimates without manual tuning.

Contribution

It proposes a cross-regularized uncertainty learning method that optimizes uncertainty parameters during training, enhancing calibration and adaptability in neural PDE models.

Findings

Learned uncertainty better calibrated on held-out data

Uncertainty concentrates in high-error regions

Framework applicable to Fourier Neural Operators

Abstract

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.