Learning OOD Robust Neural Operator with Risk-Averse Stochastic Optimization

Existing work in physical-informed machine learning (PIML) has shown that data-driven learning of solution operators can provide a fast approximate alternative to classical numerical ordinary/partial differential equations (ODEs/PDEs) solvers. Of these, Neural Operators (NOs) have emerged as particularly promising. However, a key challenge in the field of NOs lies in developing methods that can effectively handle out-of-distribution (OOD) forecasting problems. Such problems involve the ability to adaptively learn from observations of the same dynamical system governed by ODEs/PDEs, where the underlying parameters are unknown and vary across instances. These tasks further require precise predictions even when faced with initial conditions and PDEs/ODEs parameters outside the training distribution. In this study, we consider the problem of training models in a risk-reverse manner. We introduce a risk-aware framework aimed at enhancing the OOD robustness of NOs by stochastically optimizing the conditional value-at-risk (CVAR) of a loss distribution. Through experiments on different distinct OOD tasks, our approach demonstrates a significant performance improvement over existing advanced NOs.