🤖 AI Summary
This work addresses the challenge of obtaining well-calibrated uncertainty estimates from neural PDE surrogates under limited or partially observed data, where existing methods often fail to adaptively quantify uncertainty. The authors propose a cross-regularized uncertainty learning framework that jointly optimizes the predictor and low-dimensional uncertainty control parameters during training, enabling calibration-aware uncertainty modeling without post-processing or fixed noise assumptions. The approach supports learning continuous noise levels at multiple model components—including output layers, hidden features, or operator-specific elements such as spectral modes—by integrating a Fourier Neural Operator (FNO) with a gradient routing mechanism. Prediction accuracy and uncertainty calibration are optimized separately on the training and regularization sets. Evaluated on the APEBench benchmark, the method consistently achieves superior calibration across varying observation ratios and training scales, with uncertainty fields effectively highlighting high-error regions.
📝 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.