Neural combinatorial wavelet neural operator for catastrophic forgetting free in-context operator learning of multiple partial differential equations

Abstract

Machine learning has witnessed substantial growth in recent years, leading to the development of advanced deep learning models crafted to address a wide range of real-world challenges spanning various domains, including the acceleration of scientific computing. Contemporary deep learning approaches to solving partial differential equations (PDEs) involve approximating either the function mapping of a specific problem or the solution operators of a pre-defined physical system. Consequently, solving multiple PDEs representing a variety of physical systems requires training of multiple deep learning models. The creation of physics-specific models from scratch for each new physical system remains a resource-intensive undertaking, demanding considerable (i) computational time, (ii) memory resources, (iii) energy, (iv) intensive physics-specific manual tuning, and (v) large problem-specific training datasets. A more generalized machine learning-enhanced computational approach would be to learn a single unified deep learning model (commonly defined as the foundation model) instead of training multiple solvers from scratch. Besides accelerating computational simulations, such unified models will address all the above challenges. In this study, we introduce the Neural Combinatorial Wavelet Neural Operator (NCWNO) as a foundational model for scientific computing. The NCWNO leverages a gated structure that employs local wavelet integral blocks to acquire shared features across multiple physical systems, complemented by a memory-based ensembling approach among these local wavelet experts. The proposed NCWNO offers two key advantages: (i) it can simultaneously learn solution operators for multiple parametric PDEs, and (ii) with pre-training, it can be fine-tuned to new parametric PDEs with reduced training datasets and time. The proposed NCWNO is the first kernel-based foundational operator learning algorithm distinguished by its (i) integral-kernel-based learning structure, (ii) robustness against catastrophic forgetting of old PDEs, and (iii) the facilitation of knowledge transfer across dissimilar physical systems. Through an extensive set of benchmark examples, we demonstrate that the NCWNO can outperform existing multiphysics and task-specific baseline operator learning frameworks.