Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for PDEs
This addresses the reliability of scientific machine learning models for PDEs when applied to new or unseen domains, which is an incremental improvement over existing UQ methods.
The paper tackles the problem of uncertainty quantification (UQ) methods failing for out-of-domain (OOD) test inputs in neural operators for PDEs, and shows that ensembling and a proposed method, Operator-ProbConserv, improve OOD performance and satisfy physical constraints.
Existing work in scientific machine learning (SciML) has shown that data-driven learning of solution operators can provide a fast approximate alternative to classical numerical partial differential equation (PDE) solvers. Of these, Neural Operators (NOs) have emerged as particularly promising. We observe that several uncertainty quantification (UQ) methods for NOs fail for test inputs that are even moderately out-of-domain (OOD), even when the model approximates the solution well for in-domain tasks. To address this limitation, we show that ensembling several NOs can identify high-error regions and provide good uncertainty estimates that are well-correlated with prediction errors. Based on this, we propose a cost-effective alternative, DiverseNO, that mimics the properties of the ensemble by encouraging diverse predictions from its multiple heads in the last feed-forward layer. We then introduce Operator-ProbConserv, a method that uses these well-calibrated UQ estimates within the ProbConserv framework to update the model. Our empirical results show that Operator-ProbConserv enhances OOD model performance for a variety of challenging PDE problems and satisfies physical constraints such as conservation laws.