Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for PDEs

S Chandra Mouli, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, Andrew Stuart, Michael W. Mahoney, Bernie Wang
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:36372-36418, 2024.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-mouli24a, title = {Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for {PDE}s}, author = {Mouli, S Chandra and Maddix, Danielle C. and Alizadeh, Shima and Gupta, Gaurav and Stuart, Andrew and Mahoney, Michael W. and Wang, Bernie}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {36372--36418}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/mouli24a/mouli24a.pdf}, url = {https://proceedings.mlr.press/v235/mouli24a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for PDEs %A S Chandra Mouli %A Danielle C. Maddix %A Shima Alizadeh %A Gaurav Gupta %A Andrew Stuart %A Michael W. Mahoney %A Bernie Wang %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-mouli24a %I PMLR %P 36372--36418 %U https://proceedings.mlr.press/v235/mouli24a.html %V 235 %X 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.
APA
Mouli, S.C., Maddix, D.C., Alizadeh, S., Gupta, G., Stuart, A., Mahoney, M.W. & Wang, B.. (2024). Using Uncertainty Quantification to Characterize and Improve Out-of-Domain Learning for PDEs. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:36372-36418 Available from https://proceedings.mlr.press/v235/mouli24a.html.

Related Material