# Direct Learning of Calibration-Aware Uncertainty for Neural PDE Surrogates

arXiv: https://arxiv.org/abs/2602.11090 · HTML: https://mlsanity.com/paper/2602.11090/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 12 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Builds on (most cited first, 11 of 11)

- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification](https://mlsanity.com/paper/2107.07511/) (2021)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs](https://mlsanity.com/paper/2411.00180/) (2024)
- [Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction](https://mlsanity.com/paper/2402.01960/) (2024)
- [Geometric Operator Learning with Optimal Transport](https://mlsanity.com/paper/2507.20065/) (2025)
- [Light-Weight Diffusion Multiplier and Uncertainty Quantification for Fourier Neural Operators](https://mlsanity.com/paper/2508.00643/) (2025)
- [Cross-regularization: Adaptive Model Complexity through Validation Gradients](https://mlsanity.com/paper/2506.19755/) (2025)

## Abstract (arXiv)

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.